encoding_reflex.py

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