File size: 4,429 Bytes
2e0601e | 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 | """Token layout for an evidence prefix and independent decision candidates."""
from dataclasses import dataclass
from .candidate_tree import CandidateTree
from .decision_utils import MARKER, literal, locate_markers
@dataclass(frozen=True)
class DecisionLayout:
tree: CandidateTree
anchors: tuple[int, ...]
question_positions: tuple[tuple[int, ...], ...]
lexical_ids: tuple[tuple[int, ...], ...]
candidate_questions: tuple[int, ...]
question_types: tuple[int, ...]
def to_dict(self) -> dict:
return {
"prefix_length": self.tree.prefix_length,
"branch_lengths": self.tree.branch_lengths,
"anchors": self.anchors,
"question_positions": self.question_positions,
"lexical_ids": self.lexical_ids,
"candidate_questions": self.candidate_questions,
"question_types": self.question_types,
}
@classmethod
def from_dict(cls, data: dict):
return cls(
CandidateTree(data["prefix_length"], tuple(data["branch_lengths"])),
tuple(data["anchors"]), tuple(tuple(row) for row in data["question_positions"]),
tuple(tuple(row) for row in data["lexical_ids"]),
tuple(data["candidate_questions"]), tuple(data["question_types"]),
)
def validate(self, length: int):
count = len(self.tree.branch_lengths)
if self.tree.length != length or not (len(self.anchors) == len(self.lexical_ids) == len(self.candidate_questions) == count):
raise ValueError("Decision layout does not match the packed sequence")
if not self.question_types or len(self.question_positions) != len(self.question_types):
raise ValueError("Decision question metadata is incomplete")
if any(kind not in (0, 1, 2) for kind in self.question_types):
raise ValueError("Invalid decision question type")
if set(self.candidate_questions) != set(range(len(self.question_types))):
raise ValueError("Every question requires candidate branches")
start = self.tree.prefix_length
for index, size in enumerate(self.tree.branch_lengths):
if not start <= self.anchors[index] < start + size or not self.lexical_ids[index]:
raise ValueError("Invalid decision candidate metadata")
start += size
for index, positions in enumerate(self.question_positions):
branch = self.candidate_questions.index(index)
start = self.tree.prefix_length + sum(self.tree.branch_lengths[:branch])
end = start + self.tree.branch_lengths[branch]
if not positions or any(position < start or position >= end for position in positions):
raise ValueError("Question positions must belong to its first candidate branch")
def pack_candidates(tokenizer, prefix: list[int], questions: dict) -> tuple[list[int], DecisionLayout]:
ids = list(prefix)
lengths, anchors, positions, lexical, owners, kinds = [], [], [], [], [], []
marker = tokenizer.encode(MARKER.lstrip("\n"), add_special_tokens=False)
for question_index, question in enumerate(questions.values()):
kinds.append({"predicate": 0, "choice": 1, "score": 2}[question["type"]])
for candidate_index, description in enumerate(question["criteria"].values()):
text = literal({"question": question["instructions"], "criterion": description}) + MARKER
encoded = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True)
start = len(ids)
branch = encoded["input_ids"]
ids.extend(branch)
lengths.append(len(branch))
anchors.append(start + locate_markers(branch, marker, 1)[0])
owners.append(question_index)
lexical.append(tuple(tokenizer.encode(literal(description), add_special_tokens=False)))
if candidate_index == 0:
phrase = literal(question["instructions"])
begin = len('{"question": ')
end = begin + len(phrase)
positions.append(tuple(start + i for i, (a, b) in enumerate(encoded["offset_mapping"]) if a < end and b > begin))
layout = DecisionLayout(CandidateTree(len(prefix), tuple(lengths)), tuple(anchors), tuple(positions), tuple(lexical), tuple(owners), tuple(kinds))
layout.validate(len(ids))
return ids, layout
|