from transformers import AutoTokenizer, AutoModelForCausalLMfrom peft import PeftModelfrom huggingface_hub import hf_hub_downloadimport torch, torch.nn as nn, json BASE = "unsloth/gemma-3-4b-it"REPO = "hiteshluke/arbiter-4b" tok = AutoTokenizer.from_pretrained(BASE)model = AutoModelForCausalLM.from_pretrained(BASE, dtype=torch.bfloat16, device_map="auto")model = PeftModel.from_pretrained(model, REPO).eval() # Load the trained pointer headhead_path = hf_hub_download(REPO, "head.pt")head_meta = json.loads(open(hf_hub_download(REPO, "head_meta.json")).read())hidden = model.config.text_config.hidden_sizehead = nn.Linear(hidden, head_meta["num_slots"], bias=False).to( device=model.device, dtype=torch.bfloat16)head.load_state_dict({"weight": torch.load(head_path, map_location=model.device)["proj.weight"]})head.eval() @torch.no_grad()def decide(prompt: str, valid_slots: list[int]) -> int: ids = tok(prompt, return_tensors="pt", truncation=True, max_length=1024).input_ids.to(model.device) out = model(input_ids=ids, output_hidden_states=True, use_cache=False) pooled = out.hidden_states[-1][0, -1] logits = head(pooled.to(head.weight.dtype)).float().cpu().numpy() return valid_slots[int(max(range(len(valid_slots)), key=lambda i: logits[valid_slots[i]]))] # Example: BoolQ-style questionprompt = ( "State: The Eiffel Tower is in Paris, France.\n\n" "Question: Is the Eiffel Tower in France?\n\n" "Options:\nT. Yes / True\nF. No / False\n\nAnswer:")slot = decide(prompt, valid_slots=[0, 1]) # 0 = T, 1 = Fprint("Answer:", "T" if slot == 0 else "F") # -> T