import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
BASE_MODEL = "allenai/OLMo-2-1124-7B-Instruct"
ADAPTER_ID = "shreyansh12183/olmo2-7b-biomed-chem"
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_use_double_quant=True
)
base_model = AutoModelForCausalLM.from_pretrained(
BASE_MODEL,
quantization_config=bnb_config,
device_map="auto",
trust_remote_code=True
)
model = PeftModel.from_pretrained(base_model, ADAPTER_ID)
model.eval()
prompt = "Explain the core domain methodology and statutory reasoning."
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
outputs = model.generate(**inputs, max_new_tokens=300, temperature=0.3, top_p=0.9)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))