| """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 |
|
|
|
|
| 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 |
| 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: |
| 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) |
|
|
|
|
