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