File size: 2,916 Bytes
30b1727 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | """Model wrapper used at inference: PEFT-wrapped causal LM -> option logits through the restricted LM head (+ residual)."""
import os
import sys
from typing import Any, Dict, List, Optional
import torch
import torch.nn as nn
import torch.nn.functional as F
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import answer_slot as AS # noqa: E402
class DensePipeline(nn.Module):
"""PEFT-wrapped causal LM -> option logits through the restricted LM head (+ residual)."""
def __init__(self, peft_model: Any, cand_ids: List[int], softcap: Optional[float] = None):
super().__init__()
self.peft = peft_model
base = peft_model.get_base_model() if hasattr(peft_model, "get_base_model") else peft_model # merged models have no PEFT wrapper
self.backbone = base.model
hidden = base.config.get_text_config().hidden_size
w = base.lm_head.weight
self.register_buffer("cand_w", w[torch.tensor(cand_ids, device=w.device)].detach().clone(), persistent=False)
self.head = AS.SlotResidual(hidden, len(cand_ids))
self.softcap = softcap
def readout(self, hs, slot_of_opt, pos_idx, neg_idx, opt_mask):
"""Option logits from answer-slot hidden states hs (B, S, H)."""
lin = F.linear(hs, self.cand_w).float()
if self.softcap:
lin = self.softcap * torch.tanh(lin / self.softcap)
z = lin + self.head(hs).float()
zs = z.gather(1, slot_of_opt.unsqueeze(-1).expand(-1, -1, z.size(-1)))
pos = zs.gather(2, pos_idx.unsqueeze(-1)).squeeze(-1)
neg = zs.gather(2, neg_idx.clamp(min=0).unsqueeze(-1)).squeeze(-1)
return (pos - neg * (neg_idx >= 0)).masked_fill(~opt_mask, -1e4)
def forward(self, input_ids, attention_mask, slot_pos, slot_of_opt, pos_idx, neg_idx, opt_mask, return_pooled=False):
h = self.backbone(input_ids=input_ids, attention_mask=attention_mask, use_cache=False).last_hidden_state
rows = torch.arange(h.size(0), device=h.device).unsqueeze(1)
hs = h[rows, slot_pos]
logits = self.readout(hs, slot_of_opt, pos_idx, neg_idx, opt_mask)
if return_pooled: # mean answer-slot hidden state (temperature head input)
valid = (torch.arange(slot_pos.size(1), device=h.device)[None, :] < (slot_pos > 0).sum(1, keepdim=True)).unsqueeze(-1)
return logits, (hs.float() * valid).sum(1) / valid.sum(1).clamp(min=1)
return logits
def load_model(path: str, quant: str):
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
kw: Dict[str, Any] = {"dtype": torch.bfloat16, "device_map": "cuda"}
if quant == "nf4":
kw["quantization_config"] = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True)
return AutoModelForCausalLM.from_pretrained(path, **kw)
|