Views
No views yet
1import re
2import torch
3from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
4
5model_name_or_id = "OpenDFM/ChemDFM-R-14B"
6tokenizer = AutoTokenizer.from_pretrained(model_name_or_id)
7model = AutoModelForCausalLM.from_pretrained(model_name_or_id, torch_dtype=torch.float16).to("cuda")
8
9instruction = "Can you please give detailed descriptions of the molecule below?\nCl.O=C1c2c(O)cccc2-c2nn(CCNCCO)c3ccc(NCCNCCO)c1c23"
10message = [
11 {
12 "role": "system",
13 "content": "You are a helpful assistant that is good at reasoning. You always reason thoroughly before giving response. The reasoning process and answer are enclosed within <think> </think> and <ans wer> </answer> tags, respectively.\ni.e.,\n<think>\nreasoning process here\n</think>\n<answer>\nanswer here\n</answer>"
14 },
15 {
16 "role": "user",
17 "content": instruction
18 }
19]
20
21input_text = tokenizer.apply_chat_template(message, tokenize=False, add_generation_prompt=True)
22inputs = tokenizer(input_text, return_tensors="pt").to("cuda")
23generation_config = GenerationConfig(
24 do_sample=True,
25 top_k=20,
26 top_p=0.9,
27 temperature=0.9,
28 max_new_tokens=1024,
29 repetition_penalty=1.05,
30 eos_token_id=tokenizer.eos_token_id
31)
32outputs = model.generate(**inputs, generation_config=generation_config)
33
34generated_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
35input_text = tokenizer.decode(inputs["input_ids"][0], skip_special_tokens=True)
36generated_text = generated_text[len(input_text):].strip()
37print(f"{generated_text=}")
38
39thinking, answer = re.match(r'<think>(.*?)</think>\s?<answer>(.*?)</answer>', generated_text, re.DOTALL).groups()
40thinking, answer = thinking.strip(), answer.strip()
41print(f"{thinking=}")
42print(f"{answer=}")rdkit package to canonicalize the SMILES. Here is an example:1from rdkit import Chem
2def canonicalize_smiles(smiles):
3 mol = Chem.MolFromSmiles(smiles)
4 if mol is None:
5 return None
6 return Chem.MolToSmiles(mol, isomericSmiles=True, kekuleSmiles=False)1from rdkit import Chem
2def canonicalize_smiles(smiles):
3 return Chem.CanonSmiles(smiles, useChiral=True)1@misc{zhao2025chemdfmr,
2 title={ChemDFM-R: An Chemical Reasoner LLM Enhanced with Atomized Chemical Knowledge},
3 author={Zihan Zhao and Bo Chen and Ziping Wan and Lu Chen and Xuanze Lin and Shiyang Yu and Situo Zhang and Da Ma and Zichen Zhu and Danyang Zhang and Huayang Wang and Zhongyang Dai and Liyang Wen and Xin Chen and Kai Yu},
4 year={2025},
5 eprint={2507.21990},
6 archivePrefix={arXiv},
7 primaryClass={cs.CE},
8 url={https://arxiv.org/abs/2507.21990},
9}