"""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]}