import torch, torch.nn.functional as Ffrom huggingface_hub import hf_hub_downloadfrom safetensors.torch import load_filefrom transformers import AutoModel, AutoTokenizer repo = "dormantx/Qwen2.5-1.5B-Instruct-NLA-L18-ar"tok = AutoTokenizer.from_pretrained(repo)ar = AutoModel.from_pretrained(repo, torch_dtype=torch.bfloat16).cuda().eval()head = load_file(hf_hub_download(repo, "value_head.safetensors"))W, b = head["weight"].cuda(), head["bias"].cuda() if tok.pad_token is None: tok.pad_token = tok.eos_tokenenc = tok(["the outcome of a football game"], return_tensors="pt", padding=True).to("cuda")with torch.no_grad(): h = ar(**enc, output_hidden_states=True).hidden_states[18] # [B, T, 1536] last = enc["attention_mask"].sum(1) - 1 pooled = h[torch.arange(h.shape[0]), last].to(torch.float32) # last non-pad token pred = pooled @ W.T.float() + b.float() # reconstructed direction pred = F.normalize(pred, dim=-1)