from peft import PeftModelfrom transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfigimport torch bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16,) base = AutoModelForCausalLM.from_pretrained( "deepseek-ai/DeepSeek-R1-Distill-Llama-8B", quantization_config=bnb_config, device_map="auto", trust_remote_code=True,)model = PeftModel.from_pretrained( base, "Sumeetkulk/mol-grpo-optimizer-abls")model.eval() tokenizer = AutoTokenizer.from_pretrained( "Sumeetkulk/mol-grpo-optimizer-abls") system = ( "You are a medicinal chemist. Reason step by step about which structural " "changes will increase QED while keeping Tanimoto similarity between 0.3 " "and 0.7. Then output only the optimized SMILES.")seed_smiles = "CC1=CC=CC=C1Nc1ncnc2ccccc12" messages = [ {"role": "system", "content": system}, {"role": "user", "content": f"Seed: {seed_smiles}\nOptimize:"},]prompt = tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True) inputs = tokenizer(prompt, return_tensors="pt").to(model.device)with torch.no_grad(): out = model.generate( **inputs, max_new_tokens=200, temperature=0.85, do_sample=True, top_p=0.95, repetition_penalty=1.1, )completion = tokenizer.decode(out[0, inputs.input_ids.shape[1]:], skip_special_tokens=True)print(completion)# Extract: import re; re.search(r'<answer>(.*?)</answer>', completion).group(1)