File size: 8,176 Bytes
39eca46 | 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 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 | import json
import os
import torch
import torch.nn as nn
from transformers import PreTrainedModel, DebertaV2Model, AutoTokenizer
from safetensors.torch import load_file
try:
from .configuration_bert4jev import Bert4JevConfig
except Exception:
from configuration_bert4jev import Bert4JevConfig
MARKERS = ["[STATE]", "[Q]", "[OPT]"]
NOUL_OPTIONS = ("no", "yes")
OPT_SLOTS = 256
class _Collator:
def __init__(self, tok, max_state_tokens=256, max_len=512):
self.tok = tok
self.max_state = max_state_tokens
self.max_len = max_len
self.cls, self.sep = tok.cls_token_id, tok.sep_token_id
self.m_state, self.m_q, self.m_opt = tok.convert_tokens_to_ids(MARKERS)
self.OM = OPT_SLOTS
def _ids(self, text):
return self.tok(text, add_special_tokens=False)["input_ids"]
def encode_one(self, state, questions):
ids = [self.cls, self.m_state] + self._ids(state)[: self.max_state]
positions, q_positions, spans = [], [], []
for qi, q in enumerate(questions):
q_positions.append(len(ids))
t = self._ids(q["instructions"])
spans.append((-(qi + 1), len(ids) + 1, len(ids) + 1 + len(t)))
ids += [self.m_q] + t
pos = []
for oi, o in enumerate(q["options"]):
pos.append(len(ids))
t = self._ids(o)
spans.append((qi * self.OM + oi, len(ids) + 1, len(ids) + 1 + len(t)))
ids += [self.m_opt] + t
positions.append(pos)
ids.append(self.sep)
if len(ids) > self.max_len:
raise ValueError("sequence %d > max_len %d" % (len(ids), self.max_len))
return ids, positions, q_positions, spans
def __call__(self, items, device=None):
enc = [self.encode_one(s, qs) for s, qs in items]
L = max(len(e[0]) for e in enc)
Qm = max(len(e[1]) for e in enc)
Om = max(len(p) for e in enc for p in e[1])
pad = self.tok.pad_token_id
input_ids = torch.full((len(items), L), pad, dtype=torch.long)
attn = torch.zeros((len(items), L), dtype=torch.long)
opt_pos = torch.full((len(items), Qm, Om), -1, dtype=torch.long)
q_pos = torch.full((len(items), Qm), -1, dtype=torch.long)
n_slots = Qm * Om + Qm
seg = torch.full((len(items), L), n_slots, dtype=torch.long)
for b, (ids, positions, q_positions, spans) in enumerate(enc):
input_ids[b, : len(ids)] = torch.tensor(ids)
attn[b, : len(ids)] = 1
q_pos[b, : len(q_positions)] = torch.tensor(q_positions)
for qi, pos in enumerate(positions):
opt_pos[b, qi, : len(pos)] = torch.tensor(pos)
for sid, st, en in spans:
slot = (Qm * Om + (-sid - 1)) if sid < 0 else ((sid // self.OM) * self.OM + sid % self.OM)
seg[b, st:en] = slot
out = {"input_ids": input_ids, "attention_mask": attn, "opt_pos": opt_pos,
"opt_mask": opt_pos >= 0, "q_pos": q_pos, "seg": seg}
return {k: (v.to(device) if device is not None else v) for k, v in out.items()}
def _readout(kind, probs):
k = max(range(len(probs)), key=lambda i: probs[i])
conf = float(probs[k])
if kind == "choice":
return {"choice_idx": k, "confidence": conf}
if kind == "score":
return {"score": float(sum(i * p for i, p in enumerate(probs))), "confidence": conf}
return {"noul": float(probs[1])}
class Bert4JevModel(PreTrainedModel):
config_class = Bert4JevConfig
base_model_prefix = "deberta"
def __init__(self, config):
super().__init__(config)
self.deberta = DebertaV2Model(config)
h = int(getattr(config, "hidden", getattr(config, "hidden_size", 1024)))
self.head = nn.Sequential(nn.Linear(3 * h, h), nn.GELU(), nn.Linear(h, 1))
self.temperature = float(getattr(config, "temperature", 1.0))
self.pool = getattr(config, "pool", "span")
self.markers = list(getattr(config, "markers", MARKERS))
self.max_state_tokens = int(getattr(config, "max_state_tokens", 256))
self.max_len = int(getattr(config, "max_len", 512))
self._tokenizer = None
self._collator = None
def _ensure_tokenizer(self, path):
if self._tokenizer is None:
self._tokenizer = AutoTokenizer.from_pretrained(path)
self._collator = _Collator(self._tokenizer, self.max_state_tokens, self.max_len)
return self._collator
def forward(self, input_ids, attention_mask, opt_pos, opt_mask, q_pos, seg):
h = self.deberta(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
B, Qm, Om = opt_pos.shape
H = h.size(-1)
hp = h.float()
n_slots = Qm * Om + Qm
sums = hp.new_zeros((B, n_slots + 1, H)).scatter_add_(1, seg.unsqueeze(-1).expand(-1, -1, H), hp)
cnt = hp.new_zeros((B, n_slots + 1)).scatter_add_(1, seg, torch.ones_like(seg, dtype=hp.dtype)).clamp(min=1).unsqueeze(-1)
mean = sums / cnt
g = mean[:, : Qm * Om].reshape(B, Qm, Om, H)
q = mean[:, Qm * Om: Qm * Om + Qm].unsqueeze(2).expand(-1, -1, Om, -1)
g = torch.cat([q, g, q * g], dim=-1).to(self.head[0].weight.dtype)
logits = self.head(g).squeeze(-1)
return logits.masked_fill(~opt_mask, float("-inf"))
@classmethod
def _resolve_dir(cls, name, token=None):
if os.path.isdir(name):
return name
from huggingface_hub import snapshot_download
return snapshot_download(name, token=token)
@classmethod
def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
token = kwargs.pop("token", None) or kwargs.pop("use_auth_token", None)
d = cls._resolve_dir(pretrained_model_name_or_path, token=token)
config = Bert4JevConfig.from_pretrained(d)
dtype = torch.float16 if str(getattr(config, "dtype", "")).endswith("float16") else torch.float32
ai = getattr(config, "attn_implementation", "eager")
try:
backbone = DebertaV2Model.from_pretrained(d, attn_implementation=ai, dtype=dtype)
except TypeError:
backbone = DebertaV2Model.from_pretrained(d, attn_implementation=ai, torch_dtype=dtype)
model = cls(config)
model.deberta = backbone
model.head.load_state_dict(load_file(os.path.join(d, "head.safetensors")))
model.to(dtype)
model.temperature = float(getattr(config, "temperature", 1.0))
dev = kwargs.get("device") or ("cuda" if torch.cuda.is_available() else "cpu")
model.to(dev).eval()
model._ensure_tokenizer(d)
return model
@staticmethod
def _question(q):
kind = q["type"]
if kind == "noul":
return {"kind": "noul", "instructions": q["instructions"], "options": list(NOUL_OPTIONS)}
return {"kind": kind, "instructions": q["instructions"], "options": list(q["options"])}
@torch.no_grad()
def decide(self, state, questions):
if self._collator is None:
raise ValueError("tokenizer not initialized; load with from_pretrained()")
qs = [self._question(q) for q in questions]
b = self._collator([(state, qs)], self.deberta.device)
logits = self.forward(b["input_ids"], b["attention_mask"], b["opt_pos"], b["opt_mask"], b["q_pos"], b["seg"]).float()
probs = (logits / self.temperature).softmax(-1)[0]
out = []
for qi, q in enumerate(qs):
p = probs[qi, : len(q["options"])].tolist()
r = _readout(q["kind"], p)
if q["kind"] == "choice":
out.append({"choice": q["options"][r["choice_idx"]],
"probabilities": dict(zip(q["options"], p)),
"confidence": r["confidence"]})
elif q["kind"] == "score":
out.append({"score": r["score"], "probabilities": dict(zip(q["options"], p)),
"confidence": r["confidence"]})
else:
out.append({"noul": r["noul"]})
return out
|