File size: 7,241 Bytes
5e054a6 | 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 | """Request validation and the token layout used to train Reflex-1."""
from collections.abc import Mapping, Sequence
from typing import Any
import torch
def questions_of(request: Mapping[str, Any]) -> list[Mapping[str, Any]]:
if not isinstance(request, Mapping):
raise ValueError("A request must be an object")
if not isinstance(request.get("state"), str):
raise ValueError("state must be a string")
if "questions" in request:
if "question" in request or "options" in request:
raise ValueError("Use either questions or a single question/options pair")
questions = request["questions"]
if not isinstance(questions, list) or not questions:
raise ValueError("questions must be a nonempty list")
else:
questions = [request]
for question in questions:
if not isinstance(question, Mapping):
raise ValueError("Each question must be an object")
text = question.get("question")
if not isinstance(text, str) or not text.strip():
raise ValueError("question must be a nonempty string")
options = question.get("options")
if not isinstance(options, list) or not 1 <= len(options) <= 255:
raise ValueError("options must be a list of 1 to 255 strings")
if any(not isinstance(option, str) or not option.strip() for option in options):
raise ValueError("Every option must be a nonempty string")
return questions
class RequestEncoder:
"""Preserve the original two-tokenizer budget and packed-question layout.
Tokenization is performed for every request. There is no cross-request token
or embedding cache. State truncation is reported to the caller.
"""
def __init__(self, context_tokenizer, option_tokenizer):
self.context_tokenizer = context_tokenizer
self.option_tokenizer = option_tokenizer
self.max_state_tokens = 512
self.max_request_tokens = 2048
self.max_option_tokens = 512
for name in ("cls_token_id", "sep_token_id", "mask_token_id", "pad_token_id"):
if getattr(context_tokenizer, name, None) is None:
raise ValueError(f"Context tokenizer requires {name}")
if option_tokenizer.pad_token_id is None:
raise ValueError("Option tokenizer requires pad_token_id")
def encode(self, request: Mapping[str, Any]) -> dict[str, Any]:
questions = questions_of(request)
tok = self.context_tokenizer
def tokens(text):
return tok.encode(text, add_special_tokens=False)
context_prefix = tokens("CONTEXT: ")
question_prefix = tokens("\nQUESTION: ")
answer_prefix = tokens("\nANSWER: [unused0] ")
packed = []
for question in questions:
qt = tokens(question["question"])
opts = [tokens(option) for option in question["options"]]
if not qt or any(not option for option in opts):
raise ValueError("Question and option text must produce tokens")
packed.append((qt, opts))
# Keep both budget checks: changing either would change trained input truncation.
required = 2 + len(context_prefix) + sum(
len(qt) + len(question_prefix) + len(answer_prefix) + 2
+ sum(map(len, opts)) for qt, opts in packed
)
if required > self.max_request_tokens:
raise ValueError("Questions and choices exceed the 2048-token request budget")
state = tokens(request["state"])
retained = state[:min(self.max_state_tokens, self.max_request_tokens - required)]
ids = [tok.cls_token_id] + context_prefix + retained + [tok.sep_token_id]
positions, owners, groups, option_tokens = [], [], [], []
for index, ((qt, _), question) in enumerate(zip(packed, questions)):
ids += question_prefix + qt + answer_prefix
positions.append(len(ids))
ids += [tok.mask_token_id, tok.sep_token_id]
group = []
for option in question["options"]:
encoded = self.option_tokenizer.encode(option, add_special_tokens=True)
if len(encoded) > self.max_option_tokens:
raise ValueError("An option exceeds the 512-token option encoder limit")
group.append(len(option_tokens))
option_tokens.append(encoded)
owners.append(index)
groups.append(group)
if len(ids) + sum(map(len, option_tokens)) > self.max_request_tokens:
raise ValueError("Context and complete options exceed the 2048-token request budget")
return {
"ids": ids, "positions": positions, "owners": owners,
"groups": groups, "option_tokens": option_tokens,
"state_tokens": len(state), "retained_state_tokens": len(retained),
}
def batch(self, records: Sequence[dict[str, Any]], device="cpu") -> dict[str, Any]:
if not records:
raise ValueError("A batch must contain at least one request")
count = len(records)
length = max(len(record["ids"]) for record in records)
questions = max(len(record["positions"]) for record in records)
options = max(len(record["owners"]) for record in records)
ids = torch.full((count, length), self.context_tokenizer.pad_token_id, dtype=torch.long)
attention = torch.zeros_like(ids)
positions = torch.zeros((count, questions), dtype=torch.long)
owners = torch.zeros((count, options), dtype=torch.long)
valid = torch.zeros((count, options), dtype=torch.bool)
indices = torch.zeros((count, options), dtype=torch.long)
unique, lookup = [], {}
for index, record in enumerate(records):
n = len(record["ids"])
ids[index, :n] = torch.tensor(record["ids"])
attention[index, :n] = 1
positions[index, :len(record["positions"])] = torch.tensor(record["positions"])
owners[index, :len(record["owners"])] = torch.tensor(record["owners"])
valid[index, :len(record["owners"])] = True
for oi, tokens in enumerate(record["option_tokens"]):
key = tuple(tokens)
if key not in lookup:
lookup[key] = len(unique)
unique.append(tokens)
indices[index, oi] = lookup[key]
label_length = max(map(len, unique))
label_ids = torch.full(
(len(unique), label_length), self.option_tokenizer.pad_token_id, dtype=torch.long
)
label_mask = torch.zeros_like(label_ids)
for index, tokens in enumerate(unique):
label_ids[index, :len(tokens)] = torch.tensor(tokens)
label_mask[index, :len(tokens)] = 1
tensors = {
"input_ids": ids, "attention_mask": attention, "positions": positions,
"owners": owners, "option_valid": valid, "option_indices": indices,
"label_ids": label_ids, "label_mask": label_mask,
}
return {**{key: value.to(device) for key, value in tensors.items()},
"groups": [record["groups"] for record in records]}
|