language-plus-molecules/mCLM_1k-3b

1# mCLM Usage Example2 3import torch4from transformers import AutoTokenizer5from mCLM.model.qwen_based.model import Qwen2ForCausalLM6from mCLM.tokenizer.molecule_tokenizer import MoleculeTokenizer7 8# Load model and tokenizers9model = Qwen2ForCausalLM.from_pretrained(10    "YOUR_REPO_ID",11    torch_dtype=torch.bfloat16,12    device_map="auto"13)14 15tokenizer = AutoTokenizer.from_pretrained("YOUR_REPO_ID")16tokenizer.pad_token = tokenizer.eos_token17 18# Load molecule tokenizer19torch.serialization.add_safe_globals([MoleculeTokenizer])20molecule_tokenizer = torch.load("molecule_tokenizer.pth", weights_only=False)21 22# Run inference23user_input = "What is aspirin used for?"24messages = [{"role": "user", "content": user_input}]25inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")26 27outputs = model.generate(input_ids=inputs, max_new_tokens=256)28response = tokenizer.decode(outputs[0], skip_special_tokens=True)29print(response)30