from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
import torch
base = "allenai/Olmo-3-7B-Think"
bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True)
tok = AutoTokenizer.from_pretrained(base)
model = AutoModelForCausalLM.from_pretrained(base, quantization_config=bnb, device_map="auto")
model = PeftModel.from_pretrained(model, "openhubresearch/ATLAS-Taurus-Expert-7B-poc-v1")
SYS = ("You are ATLAS, a Taurus platform expert. You give precise, accurate answers "
"about the Taurus multi-agent orchestration platform: agents, runs, tools, "
"delegation, memory tiers, schedules, the admin API, and operating SOPs.")
msgs = [{"role":"system","content":SYS},
{"role":"user","content":"What determines how well an agent survives context compaction?"}]
ids = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt").to(model.device)
print(tok.decode(model.generate(ids, max_new_tokens=400, do_sample=False)[0][ids.shape[1]:], skip_special_tokens=True))