import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True)
base = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-4B",
quantization_config=bnb, device_map="auto",
trust_remote_code=True, torch_dtype=torch.bfloat16)
from modeling_plaa import FiLMLayer, PlaaCore
for i in range(16, 29):
base.model.layers[i] = FiLMLayer(base.model.layers[i])
peft = PeftModel.from_pretrained(base, "./", adapter_name="plaa")
peft.set_adapter("plaa")
import torch
ckpt = torch.load("./plaa_full.pt")
plaa_core = PlaaCore()
plaa_core.load_state_dict(ckpt["plaa_core"])
for i in range(16, 29):
peft.base_model.model.model.layers[i].scale_proj.load_state_dict(
ckpt["scale_proj"][i])
peft.base_model.model.model.layers[i].cuda()
S = plaa_core.init_state(1)
for i in range(16, 29):
peft.base_model.model.model.layers[i]._s = S
inp = tokenizer(["Hello"], return_tensors="pt").to("cuda")
out = peft.generate(**inp, max_new_tokens=50)