Ralf

class RalfConfig(PretrainedConfig):
    """Configuration carrying RALF architecture and tokenizer metadata.

    Args:
        dataset_name: Poster dataset key, usually `cgl` or `pku_posterlayout`.
        task: Canonical condition type or checkpoint task alias.
        id2label: Dataset-local label vocabulary persisted with the checkpoint.
        max_seq_length: Maximum number of layout elements.
        num_bin: Number of linear geometry bins per variable.
        var_order: Token variable order.
        special_tokens: Special tokens stored after label and geometry tokens.
        geo_quantization: Geometry quantizer name. The converted package supports
            `linear`; conversion records other values for audit.
        is_loc_vocab_shared: Whether geometry variables share one token range.
        d_model: Image encoder hidden size.
        decoder_d_model: Decoder hidden size.
        encoder_layers: Number of image encoder transformer layers.
        decoder_layers: Number of decoder layers.
        num_attention_heads: Number of attention heads.
        dropout: Dropout probability.
        retrieval_backbone: Retrieval backbone name.
        top_k: Number of retrieved examples expected by the checkpoint.
        use_reference_image: Whether retrieved reference images participate in fusion.
        layout_backbone: Layout encoder name.
        freeze_layout_encoder: Whether the layout encoder was frozen.
        fusion: Retrieval fusion variant.
        use_flag_embedding: Whether task flag embeddings are enabled.
        use_multitask: Whether checkpoint was trained as multitask.
        global_task_embedding: Whether global task embedding is enabled.
        relation_size: Maximum number of relation constraints.
        image_channels: Number of image channels consumed by the model.
        image_size: Optional `(height, width)` resize target.
        sort_order: Processor sort order metadata.
        retrieval_metadata: Retrieval-cache metadata for conversion/parity.
        original_config: Original config serialized as plain data.
        kwargs: Extra `PretrainedConfig` keyword arguments.
    """

    model_type = "ralf"

    def __init__(
        self,
        dataset_name: RalfDatasetName = "cgl",
        task: RalfConfigTaskName = "unconditional",
        id2label: Mapping[int | str, str] | None = None,
        max_seq_length: int = 10,
        num_bin: int = 128,
        var_order: Sequence[RalfLayoutVariable] = DEFAULT_VAR_ORDER,
        special_tokens: Sequence[str] = DEFAULT_SPECIAL_TOKENS,
        geo_quantization: str = "linear",
        is_loc_vocab_shared: bool = False,
        d_model: int = 256,
        decoder_d_model: int = 256,
        encoder_layers: int = 6,
        decoder_layers: int = 6,
        num_attention_heads: int = 8,
        dropout: float = 0.1,
        retrieval_backbone: str = "dreamsim",
        saliency_k: int | str = "None",
        top_k: int = 16,
        use_reference_image: bool = False,
        layout_backbone: str = "feature_extractor",
        freeze_layout_encoder: bool = True,
        fusion: str = "concat_cross_attention",
        use_flag_embedding: bool = True,
        use_multitask: bool = False,
        global_task_embedding: bool = False,
        relation_size: int = 10,
        image_channels: int = 4,
        image_size: tuple[int, int] | list[int] | None = None,
        sort_order: Sequence[str] = ("label", "lexicographic"),
        retrieval_metadata: RalfConfigMetadata | None = None,
        original_config: RalfConfigMetadata | None = None,
        original_hydra_config: RalfConfigMetadata | None = None,
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Initialize configuration values."""
        labels = (
            id2label_for_dataset(dataset_name)
            if id2label is None
            else {int(k): str(v) for k, v in id2label.items()}
        )

        self.dataset_name = dataset_name
        self.task = task
        self.id2label = labels
        self.max_seq_length = int(max_seq_length)
        self.num_bin = int(num_bin)
        self.var_order = tuple(var_order)
        self.special_tokens = tuple(special_tokens)
        self.geo_quantization = geo_quantization
        self.is_loc_vocab_shared = bool(is_loc_vocab_shared)

        self.d_model = int(d_model)
        self.decoder_d_model = int(decoder_d_model)
        self.encoder_layers = int(encoder_layers)
        self.decoder_layers = int(decoder_layers)
        self.num_attention_heads = int(num_attention_heads)
        self.dropout = float(dropout)

        self.retrieval_backbone = retrieval_backbone
        self.saliency_k = saliency_k
        self.top_k = int(top_k)
        self.use_reference_image = bool(use_reference_image)
        self.layout_backbone = layout_backbone
        self.freeze_layout_encoder = bool(freeze_layout_encoder)

        self.fusion = fusion
        self.use_flag_embedding = bool(use_flag_embedding)
        self.use_multitask = bool(use_multitask)
        self.global_task_embedding = bool(global_task_embedding)

        self.relation_size = int(relation_size)
        self.image_channels = int(image_channels)
        self.image_size = tuple(image_size) if image_size is not None else None
        self.sort_order = tuple(sort_order)
        self.retrieval_metadata = dict(retrieval_metadata or {})
        self.original_config = dict(original_config or original_hydra_config or {})
        self.original_hydra_config = self.original_config

        pad_token_id = self.special_token_id("pad")
        bos_token_id = self.special_token_id("bos")
        eos_token_id = self.special_token_id("eos")
        kwargs.pop("model_type", None)
        kwargs.pop("id2label", None)
        kwargs.pop("label2id", None)
        kwargs.pop("pad_token_id", None)
        kwargs.pop("bos_token_id", None)
        kwargs.pop("eos_token_id", None)

        super().__init__(
            id2label=self.id2label,
            label2id={label: idx for idx, label in self.id2label.items()},
        )
        self.pad_token_id = pad_token_id
        self.bos_token_id = bos_token_id
        self.eos_token_id = eos_token_id
        for key, value in kwargs.items():
            setattr(self, key, value)

    @property
    def num_bbox_tokens(self) -> int:
        """Return the number of geometry tokens."""
        return self.num_bin if self.is_loc_vocab_shared else self.num_bin * 4

    @property
    def vocab_size(self) -> int:
        """Return total autoregressive token vocabulary size."""
        return self.num_labels + self.num_bbox_tokens + len(self.special_tokens)

    @property
    def max_token_length(self) -> int:
        """Return maximum generated token length excluding BOS."""
        return self.max_seq_length * len(self.var_order)

    def special_token_id(self, name: str) -> int:
        """Return the numeric id for a special token name.

        Args:
            name: Special token name without brackets.

        Returns:
            Numeric token id.

        Raises:
            ValueError: If the token is absent from this config.
        """
        if name not in self.special_tokens:
            raise ValueError(f"Unknown special token: {name}")

        return self.num_labels + self.num_bbox_tokens + self.special_tokens.index(name)

    def bbox_token_offset(self, key: RalfLayoutVariable) -> int:
        """Return the first token id for a geometry variable."""
        if key == "label":
            return 0
        if self.is_loc_vocab_shared:
            return self.num_labels
        return self.num_labels + GEOMETRY_KEYS.index(key) * self.num_bin

class RalfImageProcessor(BaseImageProcessor):
    """Prepare RGB poster images and one-channel saliency tensors.

    Args:
        image_size: Optional `(height, width)` resize target.

    Examples:
        >>> processor = RalfImageProcessor(image_size=(8, 8))
        >>> out = processor.preprocess([torch.zeros(3, 8, 8)])
        >>> tuple(out["pixel_values"].shape)
        (1, 3, 8, 8)
    """

    model_input_names = ["pixel_values", "saliency"]

    def __init__(
        self,
        image_size: tuple[int, int] | None = None,
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Initialize image resize metadata."""
        super().__init__(**kwargs)  # ty: ignore[invalid-argument-type]
        self.image_size = tuple(image_size) if image_size is not None else None

    def preprocess(
        self,
        images: ImageInput | Sequence[ImageInput] | None,
        saliency: ImageInput | Sequence[ImageInput] | None = None,
        return_tensors: Literal["pt"] = "pt",
        **kwargs: str | int | float | bool | None,
    ) -> BatchFeature:
        """Convert images and saliency maps to tensors.

        Args:
            images: RGB images as PIL, NumPy, or torch tensors.
            saliency: Optional single-channel saliency maps.
            return_tensors: Tensor return format. Only `pt` is supported.
            kwargs: Reserved processor arguments.

        Returns:
            BatchFeature with `pixel_values` and `saliency`.

        Raises:
            ValueError: If `return_tensors` is not `pt`.
        """
        _ = kwargs
        if return_tensors != "pt":
            raise ValueError("RalfImageProcessor supports return_tensors='pt' only")

        image_items = (
            _as_list(images) if images is not None else [torch.zeros(3, 64, 64)]
        )
        pixel_values = torch.stack(
            [_image_to_tensor(item, channels=3) for item in image_items]
        )
        if self.image_size is not None:
            pixel_values = torch.nn.functional.interpolate(
                pixel_values,
                size=self.image_size,
                mode="bilinear",
                align_corners=False,
            )
        if saliency is None:
            saliency_values = torch.zeros(
                pixel_values.size(0),
                1,
                pixel_values.size(2),
                pixel_values.size(3),
                dtype=pixel_values.dtype,
            )
        else:
            saliency_items = _as_list(saliency)
            if len(saliency_items) == 1 and pixel_values.size(0) > 1:
                saliency_items = saliency_items * pixel_values.size(0)
            saliency_values = torch.stack(
                [_image_to_tensor(item, channels=1) for item in saliency_items]
            )
            if self.image_size is not None:
                saliency_values = torch.nn.functional.interpolate(
                    saliency_values,
                    size=self.image_size,
                    mode="bilinear",
                    align_corners=False,
                )
        return BatchFeature(
            {"pixel_values": pixel_values, "saliency": saliency_values},
            tensor_type=return_tensors,
        )

    def to_dict(self) -> dict[str, RalfImageProcessorConfigValue]:
        """Serialize image processor metadata."""
        data = cast(dict[str, RalfImageProcessorConfigValue], super().to_dict())
        data["image_size"] = self.image_size
        return data

class RalfForConditionalLayoutGeneration(PreTrainedModel):
    """Standalone `PreTrainedModel` for RALF autoregressive decoding."""

    config_class = RalfConfig
    base_model_prefix = "ralf"
    main_input_name = "input_ids"
    _tied_weights_keys: dict[str, str] = {}

    flag_img: Int[torch.Tensor, "1"]
    flag_user_const: Int[torch.Tensor, "1"]

    def __init__(self, config: RalfConfig) -> None:
        """Initialize a local module tree matching original RALF checkpoint keys."""
        super().__init__(config)
        self.tokenizer = RalfTokenizerView(config)
        self.dataset_name = config.dataset_name
        self.d_model = config.d_model
        self.max_seq_length = config.max_seq_length

        self.use_reference_image = config.use_reference_image
        self.layout_backbone = config.layout_backbone
        self.top_k = config.top_k
        self.weight_init = True

        self.retrieval_backbone = config.retrieval_backbone
        self.random_retrieval = False
        self.saliency_k = str(config.saliency_k)

        self.num_layers = config.encoder_layers
        self.nhead = config.num_attention_heads
        self.dropout = config.dropout

        self.encoder = ResnetFeatureExtractor(
            backbone="resnet50", d_model=config.d_model, head="transformer"
        )
        self.pos_emb_2d = PositionEmbeddingSine(config.d_model, normalize=True)
        self.dim_feedforward = 4 * config.d_model
        self.transformer_encoder = nn.TransformerEncoder(
            encoder_layer=nn.TransformerEncoderLayer(
                d_model=config.d_model,
                nhead=config.num_attention_heads,
                batch_first=True,
                dropout=config.dropout,
                norm_first=True,
                dim_feedforward=self.dim_feedforward,
            ),
            num_layers=config.encoder_layers,
        )
        self.decoder = BaseDecoder(
            d_label=self.tokenizer.N_total,
            d_model=config.decoder_d_model,
            num_layers=config.decoder_layers,
            nhead=config.num_attention_heads,
            pos_emb="layout",
            dim_feedforward=self.dim_feedforward,
        )
        self.loss_fn_ce = nn.CrossEntropyLoss(
            label_smoothing=0.1, ignore_index=self.tokenizer.name_to_id("pad")
        )
        self.layout_encoer = FIDNetFeatureExtractor(
            num_label=self.tokenizer.N_label,
            d_model=256,
            nhead=4,
            num_layers=4,
            max_bbox=config.max_seq_length,
        )
        self.layout_encoer.enc_transformer.token.requires_grad = False
        for parameter in self.layout_encoer.parameters():
            parameter.requires_grad = False
        self.pos_emb_1d = PositionalEncoding1d(
            d_model=config.d_model,
            max_len=5000 if not config.use_reference_image else 10000,
        )
        self.layout_adapter = FeedForward(
            dim=256, hidden_dim=4 * config.d_model, output_dim=config.d_model
        )
        self.head = FeedForward(dim=config.d_model, hidden_dim=4 * config.d_model)
        self.auxilary_task = self._canonical_to_task_name(config.task)
        self.use_multitask = config.use_multitask
        self.global_task_embedding = config.global_task_embedding
        self.preprocessor = RalfTaskPreprocessor(
            tokenizer=self.tokenizer,
            task=self.auxilary_task,
            global_task_embedding=config.global_task_embedding,
        )
        self.user_const_encoder = UserConstraintTransformerEncoder(
            d_model=config.d_model,
            nhead=config.num_attention_heads,
            num_layers=config.encoder_layers,
            d_label=self.preprocessor.N_total,
            dim_feedforward=self.dim_feedforward,
        )
        self.use_flag_embedding = config.use_flag_embedding
        if self.use_flag_embedding:
            self.task_emb = nn.Embedding(2, 1)
            nn.init.normal_(self.task_emb.weight, mean=0.0, std=0.02)
            self.register_buffer("flag_img", torch.zeros(1).long())
            self.register_buffer("flag_user_const", torch.ones(1).long())
        self.attn = Attention(
            config.d_model, config.d_model, heads=8, dim_head=64, dropout=0.0
        )
        self.all_tied_weights_keys = dict(self._tied_weights_keys)

    @staticmethod
    def _canonical_to_task_name(task: RalfConfigTaskName | str) -> RalfTaskName:
        if task not in TASK_BY_CONDITION:
            raise ValueError(f"Unsupported RALF task or condition: {task}")

        return TASK_BY_CONDITION[cast(RalfConfigTaskName, task)]

    def _default_retrieved(
        self, batch_size: int, device: torch.device, dtype: torch.dtype
    ) -> dict[str, Shaped[torch.Tensor, ...]]:
        shape = (batch_size, self.config.top_k, self.config.max_seq_length)
        image = torch.zeros(
            batch_size, self.config.top_k, 4, 64, 64, device=device, dtype=dtype
        )
        return {
            "image": image,
            "center_x": torch.zeros(shape, device=device, dtype=dtype),
            "center_y": torch.zeros(shape, device=device, dtype=dtype),
            "width": torch.zeros(shape, device=device, dtype=dtype),
            "height": torch.zeros(shape, device=device, dtype=dtype),
            "label": torch.zeros(shape, device=device, dtype=torch.long),
            "mask": torch.zeros(shape, device=device, dtype=torch.bool),
        }

    def _encode_into_memory(
        self,
        inputs: Mapping[
            str, Shaped[torch.Tensor, ...] | Mapping[str, Shaped[torch.Tensor, ...]]
        ],
    ) -> dict[str, Float[torch.Tensor, "batch memory_tokens channels"]]:
        image = cast(Tensor, inputs["image"])
        retrieved = cast(Mapping[str, Shaped[torch.Tensor, "..."]], inputs["retrieved"])
        input_img_feature = self.encoder(image)
        input_img_feature = self.pos_emb_2d(input_img_feature)
        image_memory = self.transformer_encoder(input_img_feature)
        ref_layouts = _extract_retrieved_features(
            retrieved_samples=retrieved,
            top_k=self.top_k,
            layout_encoder=self.layout_encoer,
            layout_adapter=self.layout_adapter,
            pos_emb_1d=self.pos_emb_1d,
        )
        memory_ca = self.attn(image_memory, ref_layouts)
        img_retrieved_layout_memory = self.head(
            torch.cat([image_memory, memory_ca, ref_layouts], dim=1)
        )
        if self.global_task_embedding:
            task_token = self.preprocessor.get_token(
                self.preprocessor.TASK, img_retrieved_layout_memory.size(0)
            ).type_as(cast(Tensor, inputs["seq_layout_const"]))
        else:
            task_token = None
        user_const_feature = self.user_const_encoder(
            src=cast(Tensor, inputs["seq_layout_const"]),
            src_key_padding_mask=cast(Tensor, inputs["seq_layout_const_pad_mask"]),
            task_token=task_token,
        )
        if self.use_flag_embedding:
            img_retrieved_layout_memory = img_retrieved_layout_memory + self.task_emb(
                self.flag_img
            )
            user_const_feature = user_const_feature + self.task_emb(
                self.flag_user_const
            )
        return {
            "memory": torch.cat(
                [img_retrieved_layout_memory, user_const_feature], dim=1
            )
        }

    def _prepare_conditional_inputs(
        self,
        *,
        pixel_values: Float[torch.Tensor, "batch channels height width"] | None,
        saliency: Float[torch.Tensor, "batch 1 height width"] | None,
        retrieved: RalfRetrievedBatch | None,
        batch_size: int,
        condition_type: RalfConfigTaskName | None = None,
        constraint_input_ids: Int[torch.Tensor, "batch tokens"] | None = None,
        constraint_mask: Bool[torch.Tensor, "batch tokens"] | None = None,
        constraint_element_mask: Bool[torch.Tensor, "batch elements"] | None = None,
        relationship_table: RalfRelationshipTable | None = None,
        sample_ids: Int[torch.Tensor, "batch"]
        | Sequence[int | str]
        | int
        | str
        | None = None,
    ) -> dict[str, Shaped[torch.Tensor, ...] | Mapping[str, Shaped[torch.Tensor, ...]]]:
        device = next(self.parameters()).device
        dtype = next(self.parameters()).dtype
        if pixel_values is None:
            pixel_values = torch.zeros(
                batch_size, 3, 64, 64, device=device, dtype=dtype
            )
        if saliency is None:
            saliency = torch.zeros(
                pixel_values.size(0),
                1,
                pixel_values.size(2),
                pixel_values.size(3),
                device=pixel_values.device,
                dtype=pixel_values.dtype,
            )
        if pixel_values.size(-1) < 64 or pixel_values.size(-2) < 64:
            pixel_values = F.interpolate(
                pixel_values, size=(64, 64), mode="bilinear", align_corners=False
            )
            saliency = F.interpolate(
                saliency, size=(64, 64), mode="bilinear", align_corners=False
            )
        if pixel_values.size(0) == 1 and batch_size > 1:
            pixel_values = pixel_values.expand(batch_size, -1, -1, -1)
            saliency = saliency.expand(batch_size, -1, -1, -1)
        image = torch.cat([pixel_values, saliency], dim=1).to(
            device=device, dtype=dtype
        )
        retrieved_dict = (
            self._default_retrieved(image.size(0), device, dtype)
            if retrieved is None
            else {
                key: value.to(device=device, dtype=dtype)
                if value.is_floating_point()
                else value.to(device=device)
                for key, value in retrieved_batch_to_model_inputs(retrieved).items()
            }
        )
        if retrieved_dict["image"].size(2) == 3:
            retrieved_dict["image"] = torch.cat(
                [retrieved_dict["image"], retrieved_dict["saliency"]], dim=2
            )
        task = self._canonical_to_task_name(condition_type or self.auxilary_task)
        preprocessor = (
            self.preprocessor
            if task == self.auxilary_task and relationship_table is None
            else RalfTaskPreprocessor(
                tokenizer=self.tokenizer,
                task=task,
                global_task_embedding=self.global_task_embedding,
                relationship_table=relationship_table if task == "relation" else None,
                relation_size=self.config.relation_size,
            )
        )
        cond = RalfConditionalInputs(
            image=image,
            retrieved=retrieved_dict,
            seq=constraint_input_ids,
            mask=constraint_mask,
            element_mask=constraint_element_mask,
            task=task,
            id=sample_ids,
        )
        seq_constraints = preprocessor(cond)
        return {
            "image": image,
            "retrieved": retrieved_dict,
            "seq_layout_const": seq_constraints["seq"],
            "seq_layout_const_pad_mask": seq_constraints["pad_mask"],
        }

    def _prepare_unconditional_inputs(
        self,
        *,
        pixel_values: Float[torch.Tensor, "batch channels height width"] | None,
        saliency: Float[torch.Tensor, "batch 1 height width"] | None,
        retrieved: RalfRetrievedBatch | None,
        batch_size: int,
    ) -> dict[str, Shaped[torch.Tensor, ...] | Mapping[str, Shaped[torch.Tensor, ...]]]:
        return self._prepare_conditional_inputs(
            pixel_values=pixel_values,
            saliency=saliency,
            retrieved=retrieved,
            batch_size=batch_size,
            condition_type="uncond",
        )

    def forward(
        self,
        input_ids: Int[torch.Tensor, "batch tokens"] | None = None,
        pixel_values: Float[torch.Tensor, "batch channels height width"] | None = None,
        saliency: Float[torch.Tensor, "batch 1 height width"] | None = None,
        attention_mask: Bool[torch.Tensor, "batch tokens"] | None = None,
        labels: Int[torch.Tensor, "batch tokens"] | None = None,
        retrieved: RalfRetrievedBatch | None = None,
        condition_type: RalfConfigTaskName | None = None,
        constraint_element_mask: Bool[torch.Tensor, "batch elements"] | None = None,
        return_dict: bool | None = None,
        **kwargs: str | float | bool | None,
    ) -> CausalLMOutput | tuple[Float[torch.Tensor, ...], ...]:
        """Run teacher-forced token prediction using the local RALF port."""
        relationship_table = cast(
            RalfRelationshipTable | None, kwargs.get("relationship_table")
        )
        sample_ids = cast(
            Int[torch.Tensor, "batch"] | Sequence[int | str] | int | str | None,
            kwargs.get("sample_ids"),
        )
        if input_ids is None:
            raise ValueError("input_ids is required")

        encoder_inputs = self._prepare_conditional_inputs(
            pixel_values=pixel_values,
            saliency=saliency,
            retrieved=retrieved,
            batch_size=input_ids.size(0),
            condition_type=condition_type,
            constraint_input_ids=input_ids,
            constraint_mask=attention_mask,
            constraint_element_mask=constraint_element_mask,
            relationship_table=relationship_table,
            sample_ids=sample_ids,
        )
        encoded_feat = self._encode_into_memory(encoder_inputs)
        logits = self.decoder(
            tgt=input_ids,
            tgt_key_padding_mask=None
            if attention_mask is None
            else ~attention_mask.bool(),
            is_causal=True,
            **encoded_feat,
        )
        loss = None
        if labels is not None:
            targets = labels.clone()
            targets[targets == self.config.pad_token_id] = -100
            loss = F.cross_entropy(
                logits.reshape(-1, logits.size(-1)),
                targets.reshape(-1),
                ignore_index=-100,
            )
        if return_dict is False:
            return (logits,) if loss is None else (loss, logits)
        return CausalLMOutput(loss=cast(torch.FloatTensor | None, loss), logits=logits)

    @torch.no_grad()
    def _generate_sequences(
        self,
        input_ids: Int[torch.Tensor, "batch tokens"],
        pixel_values: Float[torch.Tensor, "batch channels height width"] | None = None,
        saliency: Float[torch.Tensor, "batch 1 height width"] | None = None,
        attention_mask: Bool[torch.Tensor, "batch tokens"] | None = None,
        *,
        max_length: int | None = None,
        temperature: float = 1.0,
        top_k: int | None = None,
        generator: torch.Generator | None = None,
        token_mask: Bool[torch.Tensor, "tokens vocab"] | None = None,
        retrieved: RalfRetrievedBatch | None = None,
        condition_type: RalfConfigTaskName | None = None,
        constraint_input_ids: Int[torch.Tensor, "batch tokens"] | None = None,
        constraint_mask: Bool[torch.Tensor, "batch tokens"] | None = None,
        constraint_element_mask: Bool[torch.Tensor, "batch elements"] | None = None,
        relationship_table: RalfRelationshipTable | None = None,
        sample_ids: Int[torch.Tensor, "batch"]
        | Sequence[int | str]
        | int
        | str
        | None = None,
    ) -> Int[torch.Tensor, "batch tokens"]:
        """Run the RALF autoregressive token loop used by `RalfPipeline`."""
        _ = attention_mask
        was_training = self.training
        self.eval()
        task = self._canonical_to_task_name(condition_type or self.auxilary_task)
        generated = input_ids[:, :1].clone()
        start_step = 0
        if task == "partial":
            condition_seq = (
                constraint_input_ids if constraint_input_ids is not None else input_ids
            )
            prefix = condition_seq[:, 1 : 1 + len(self.config.var_order)]
            generated = torch.cat([generated, prefix.to(generated.device)], dim=1)
            start_step = len(self.config.var_order)
        max_length = max_length or self.config.max_token_length
        encoder_inputs = self._prepare_conditional_inputs(
            pixel_values=pixel_values,
            saliency=saliency,
            retrieved=retrieved,
            batch_size=input_ids.size(0),
            condition_type=task,
            constraint_input_ids=constraint_input_ids,
            constraint_mask=constraint_mask,
            constraint_element_mask=constraint_element_mask,
            relationship_table=relationship_table,
            sample_ids=sample_ids,
        )
        encoded_feat = self._encode_into_memory(encoder_inputs)
        try:
            for step in range(start_step, max_length):
                logits = self.decoder(
                    tgt=generated,
                    tgt_key_padding_mask=generated.eq(self.config.pad_token_id),
                    is_causal=True,
                    **encoded_feat,
                )
                next_logits = logits[:, step : step + 1]
                next_logits = rearrange(next_logits, "b 1 c -> b c") / temperature
                if token_mask is not None and step < token_mask.size(0):
                    next_logits = next_logits.masked_fill(
                        ~token_mask[step].to(next_logits.device), -math.inf
                    )
                next_logits = _apply_decode_space_restriction(
                    task=task,
                    step=step,
                    condition=constraint_input_ids,
                    logits=next_logits,
                    pad_id=self.config.pad_token_id,
                    eos_id=self.config.eos_token_id,
                    max_length=self.config.max_token_length,
                )
                if top_k is not None and top_k > 0 and top_k < next_logits.size(-1):
                    values = torch.topk(next_logits, top_k).values
                    next_logits = next_logits.masked_fill(
                        next_logits < values[:, [-1]], -math.inf
                    )
                probs = F.softmax(next_logits, dim=-1)
                next_token = torch.multinomial(
                    probs, num_samples=1, generator=generator
                )
                generated = torch.cat([generated, next_token], dim=1)
        finally:
            if was_training:
                self.train()
        return generated