import torchfrom transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfigfrom peft import PeftModel # Load base modelbnb = 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) # Inject FiLM layersfrom modeling_plaa import FiLMLayer, PlaaCorefor i in range(16, 29): base.model.layers[i] = FiLMLayer(base.model.layers[i]) # Load PEFT adapterpeft = PeftModel.from_pretrained(base, "./", adapter_name="plaa")peft.set_adapter("plaa") # Load FiLM weightsimport torchckpt = 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() # InferenceS = plaa_core.init_state(1)for i in range(16, 29): peft.base_model.model.model.layers[i]._s = Sinp = tokenizer(["Hello"], return_tensors="pt").to("cuda")out = peft.generate(**inp, max_new_tokens=50)