File size: 6,748 Bytes
8a48617 | 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 | """One joint row per question. Metadata IDs and gold labels never enter tokens."""
import math
import torch
KINDS = ("choice", "noul", "score")
DEFAULT_NO = "No. The statement or question is not satisfied."
DEFAULT_YES = "Yes. The statement or question is satisfied."
def nonempty(value, name):
if not isinstance(value, str) or not value.strip():
raise ValueError(f"{name} must be a nonempty string")
return value
def view(record, labeled=False):
q = record["question"]
kind = str(q["type"]).lower()
if kind not in KINDS:
raise ValueError("Unsupported question type")
nonempty(record["state_text"], "state_text")
nonempty(q["text"], "question.text")
if kind == "noul":
options = [{"id": "no", "text": q.get("false_criterion", DEFAULT_NO)},
{"id": "yes", "text": q.get("true_criterion", DEFAULT_YES)}]
else:
options = q["options" if kind == "choice" else "levels"]
if not isinstance(options, list) or not 2 <= len(options) <= 255:
raise ValueError("Require 2..255 complete candidate descriptions")
for o in options:
nonempty(o["id"], "candidate.id")
nonempty(o["text"], "candidate.text")
if len({o["id"] for o in options}) != len(options):
raise ValueError("Duplicate candidate IDs")
values = [float(o["value"]) for o in options] if kind == "score" else list(range(len(options)))
if kind == "score" and (not all(math.isfinite(v) for v in values)
or any(a >= b for a, b in zip(values, values[1:]))):
raise ValueError("Score level values must be finite and strictly increasing")
target = None
if labeled:
gold = record["target"]
if kind == "noul":
yes = float(gold["probability"])
target = [1 - yes, yes]
elif "probabilities" in gold:
target = [float(x) for x in gold["probabilities"]]
elif kind == "choice" and "choice_id" in gold:
if gold["choice_id"] not in {o["id"] for o in options}:
raise ValueError("Unknown gold choice_id")
target = [float(o["id"] == gold["choice_id"]) for o in options]
else:
raise ValueError("Score requires an explicit target probability distribution")
if (len(target) != len(options) or not all(math.isfinite(x) and 0 <= x <= 1 for x in target)
or abs(sum(target) - 1) > 1e-6):
raise ValueError("Invalid target probability distribution")
return kind, options, values, target
class MarkerCollator:
def __init__(self, tokenizer, max_length=8192, state_truncation="error"):
if isinstance(max_length, bool) or not isinstance(max_length, int) or max_length < 8:
raise ValueError("Invalid max_length")
if state_truncation not in ("error", "right"):
raise ValueError("state_truncation must be error or right")
self.tokenizer, self.max_length = tokenizer, max_length
self.state_truncation = state_truncation
self.bos = tokenizer.cls_token_id if tokenizer.cls_token_id is not None else tokenizer.bos_token_id
self.sep = tokenizer.sep_token_id if tokenizer.sep_token_id is not None else tokenizer.eos_token_id
self.pad, self.marker = tokenizer.pad_token_id, tokenizer.mask_token_id
if any(v is None for v in (self.bos, self.sep, self.pad, self.marker)):
raise ValueError("Tokenizer must provide existing BOS/CLS, SEP/EOS, PAD and MASK IDs")
def tokens(self, text):
ids = self.tokenizer(text, add_special_tokens=False, truncation=False)["input_ids"]
if not ids:
raise ValueError("Empty token span")
return ids
def encode(self, record, labeled=False):
kind, options, values, target = view(record, labeled)
question = self.tokens(f"{kind} question: {record['question']['text']}")
ids = [self.bos] + question + [self.sep]
positions, spans = [], []
for j, option in enumerate(options):
# Ordinal rank is meaningful only for Score; opaque external IDs never enter text.
text = f"level {j}: {option['text']}" if kind == "score" else option["text"]
desc = self.tokens(text)
positions.append(len(ids))
ids.append(self.marker)
spans.append([len(ids), len(ids) + len(desc)])
ids.extend(desc + [self.sep])
state = self.tokens(record["state_text"])
room = self.max_length - len(ids) - 1
if room < 1:
raise ValueError("Complete question and candidates leave no room for state")
if len(state) > room and self.state_truncation == "error":
raise ValueError(f"Record {record.get('id')} exceeds {self.max_length} tokens; no implicit truncation")
kept = state[:room]
state_span = [len(ids), len(ids) + len(kept)]
ids.extend(kept + [self.sep])
return {"ids": ids, "positions": positions, "candidate_spans": spans,
"state_span": state_span, "kind": kind, "target": target, "values": values,
"candidate_ids": [o["id"] for o in options], "id": record.get("id"),
"question_id": record["question"].get("id"), "state_tokens_original": len(state),
"state_tokens_kept": len(kept), "input_tokens": len(ids)}
def __call__(self, records, labeled=False, device="cpu"):
encoded = [self.encode(r, labeled) for r in records]
if not encoded:
raise ValueError("Empty batch")
b, length, k = len(encoded), max(len(e["ids"]) for e in encoded), max(len(e["positions"]) for e in encoded)
ids = torch.full((b, length), self.pad, dtype=torch.long)
mask = torch.zeros((b, length), dtype=torch.bool)
positions = torch.zeros((b, k), dtype=torch.long)
valid = torch.zeros((b, k), dtype=torch.bool)
target, values = torch.zeros((b, k)), torch.zeros((b, k))
for i, e in enumerate(encoded):
n, count = len(e["ids"]), len(e["positions"])
ids[i, :n] = torch.tensor(e["ids"])
mask[i, :n] = True
positions[i, :count] = torch.tensor(e["positions"])
valid[i, :count] = True
values[i, :count] = torch.tensor(e["values"])
if labeled:
target[i, :count] = torch.tensor(e["target"])
batch = {"input_ids": ids, "attention_mask": mask, "marker_positions": positions,
"valid_candidates": valid, "targets": target, "values": values,
"kind_ids": torch.tensor([KINDS.index(e["kind"]) for e in encoded])}
return {name: value.to(device) for name, value in batch.items()}, encoded
|