from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import 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)
