Layout action

class LayoutActionConfig(PretrainedConfig):
    """Architecture and tokenizer metadata for LayoutAction checkpoints.

    Args:
        dataset_name: Dataset slug. ``rico`` and ``layout_action_rico13`` map to
            the released RICO13 label order.
        id2label: Dataset-local label mapping.
        precision: Coordinate precision; checkpoint defaults to 8 bits.
        max_elements: Maximum number of layout elements.
        block_size: GPT context length. Defaults to released max token length.
        vocab_size: Token vocabulary size. Defaults to the reference formula.
        n_layer: Number of GPT blocks.
        n_head: Attention heads.
        n_embd: Hidden size.
        embd_pdrop: Embedding dropout.
        resid_pdrop: Residual dropout.
        attn_pdrop: Attention dropout.
        default_sampling: Default pipeline sampling mode.
        default_top_k: Default top-k value.
        default_temperature: Default sampling temperature.
        original_dataset_name: Original dataset name.
        original_asset_manifest: Optional asset manifest.
        kwargs: Additional ``PretrainedConfig`` fields.

    Examples:
        >>> config = LayoutActionConfig(dataset_name="publaynet")
        >>> config.bos_token_id == config.vocab_size - 3
        True
    """

    model_type = "layout-action"

    def __init__(
        self,
        *,
        dataset_name: str = "rico13",
        id2label: Mapping[int, str] | Mapping[str, str] | None = None,
        precision: int = 8,
        max_elements: int | None = None,
        block_size: int | None = None,
        vocab_size: int | None = None,
        n_layer: int = 6,
        n_head: int = 8,
        n_embd: int = 512,
        embd_pdrop: float = 0.1,
        resid_pdrop: float = 0.1,
        attn_pdrop: float = 0.1,
        default_sampling: LayoutActionSamplingMode
        | str = LayoutActionSamplingMode.top_k,
        default_top_k: int = 5,
        default_temperature: float = 1.0,
        original_dataset_name: str | None = None,
        original_asset_manifest: Mapping[
            str, str | int | list[str] | dict[str, AssetManifestFile]
        ]
        | None = None,
        model_type: str | None = None,
        transformers_version: str | None = None,
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Initialize LayoutAction metadata and derived token ids."""
        _ = (model_type, transformers_version)
        for derived_key in (
            "bos_token_id",
            "eos_token_id",
            "pad_token_id",
            "label2id",
            "size",
            "element_token_width",
            "max_token_length",
            "no_value_token_id",
            "label_token_offset",
            "copy_token_id",
            "margin_token_id",
            "generate_token_id",
            "no_obj_token_id",
        ):
            kwargs.pop(derived_key, None)
        _ = kwargs
        super().__init__()
        dataset = normalize_vendor_dataset_name(dataset_name)
        labels = layout_action_labels(dataset)
        normalized_id2label = (
            {int(key): str(value) for key, value in id2label.items()}
            if id2label is not None
            else dict(enumerate(labels))
        )
        self.dataset_name = dataset
        self.precision = int(precision)
        self.max_elements = int(
            max_elements
            if max_elements is not None
            else max_elements_for_layout_action_dataset(dataset)
        )
        self.n_layer = int(n_layer)
        self.n_head = int(n_head)
        self.n_embd = int(n_embd)
        self.embd_pdrop = float(embd_pdrop)
        self.resid_pdrop = float(resid_pdrop)
        self.attn_pdrop = float(attn_pdrop)
        self.default_sampling = str(normalize_sampling_mode(default_sampling))
        self.default_top_k = int(default_top_k)
        self.default_temperature = float(default_temperature)
        self.original_dataset_name = original_dataset_name or dataset
        self.original_asset_manifest = dict(original_asset_manifest or {})
        self.id2label: dict[int, str] = normalized_id2label
        self.label2id: dict[str, int] = {
            value: key for key, value in self.id2label.items()
        }
        self.size: int = 2**self.precision
        self.element_token_width: int = ELEMENT_TOKEN_WIDTH
        self.max_token_length: int = self.max_elements * self.element_token_width + 2
        self.block_size = int(
            block_size if block_size is not None else self.max_token_length
        )
        self.no_value_token_id: int = self.size
        self.label_token_offset: int = self.size + 1
        self.copy_token_id: int = self.label_token_offset + len(self.id2label)
        self.margin_token_id: int = self.copy_token_id + 1
        self.generate_token_id: int = self.margin_token_id + 1
        self.no_obj_token_id: int = self.generate_token_id + 1
        resolved_vocab_size = (
            int(vocab_size)
            if vocab_size is not None
            else self.size + 1 + len(self.id2label) + 3 + 1 + self.max_elements + 3
        )
        self.vocab_size = resolved_vocab_size
        self.bos_token_id = self.vocab_size - 3
        self.eos_token_id = self.vocab_size - 2
        self.pad_token_id = self.vocab_size - 1

    def label_token_id(self, label_id: int) -> int:
        """Return the synthetic token id for a dataset-local label id."""
        id2label = cast(dict[int, str], self.id2label)
        if label_id not in id2label:
            raise ValueError(f"Unknown LayoutAction label id: {label_id}")

        return self.label_token_offset + int(label_id)

    def label_id_from_token(self, token_id: int) -> int | None:
        """Return a dataset-local label id for a label token id."""
        id2label = cast(dict[int, str], self.id2label)
        label_id = int(token_id) - self.label_token_offset
        if 0 <= label_id < len(id2label):
            return label_id
        return None

    def object_token_id(self, back_reference: int) -> int:
        """Return the token id for a previous-object back reference."""
        if not 1 <= back_reference <= self.max_elements:
            raise ValueError("back_reference must be in [1, max_elements]")

        return self.no_obj_token_id + back_reference

    def back_reference_from_token(self, token_id: int) -> int | None:
        """Return a previous-object back reference from a token id."""
        value = int(token_id) - self.no_obj_token_id
        if 1 <= value <= self.max_elements:
            return value
        return None

class LayoutActionSamplingMode(StrEnum):
    """Supported token sampling modes."""

    greedy = auto()
    multinomial = auto()
    top_k = auto()

@dataclass(frozen=True)
class StateDictKeyReport:
    """One source-to-target state-dict mapping result."""

    source_key: str
    target_key: str
    source_shape: tuple[int, ...]
    loaded: bool

@dataclass(frozen=True)
class LayoutActionSamplingConfig:
    """Sampling parameters for LayoutAction token generation.

    Args:
        mode: Greedy, multinomial, or top-k sampling.
        temperature: Positive logit temperature.
        top_k: Optional top-k crop size.

    Examples:
        >>> str(LayoutActionSamplingConfig().mode)
        'top_k'
    """

    mode: LayoutActionSamplingMode = LayoutActionSamplingMode.top_k
    temperature: float = 1.0
    top_k: int | None = 5

    @classmethod
    def from_values(
        cls,
        *,
        mode: LayoutActionSamplingMode | str,
        temperature: float,
        top_k: int | None,
    ) -> "LayoutActionSamplingConfig":
        """Build a normalized sampling config from public values."""
        return cls(
            mode=normalize_sampling_mode(mode),
            temperature=float(temperature),
            top_k=None if top_k is None else int(top_k),
        )

class LayoutActionForCausalLM(PreTrainedModel):
    """Transformers ``PreTrainedModel`` for LayoutAction token prediction.

    Args:
        config: LayoutAction architecture and vocabulary metadata.

    Examples:
        >>> config = LayoutActionConfig(n_layer=1, n_head=2, n_embd=16, max_elements=1)
        >>> model = LayoutActionForCausalLM(config)
        >>> out = model(torch.tensor([[config.bos_token_id]]))
        >>> out.logits.shape[-1] == config.vocab_size
        True
    """

    config_class = LayoutActionConfig
    base_model_prefix = "layout_action"
    main_input_name = "input_ids"
    _tied_weights_keys: dict[str, str] = {}

    def __init__(self, config: LayoutActionConfig) -> None:
        """Initialize checkpoint-compatible GPT modules."""
        super().__init__(config)
        self.tok_emb = nn.Embedding(config.vocab_size, config.n_embd)
        self.pos_emb = nn.Parameter(torch.zeros(1, config.block_size, config.n_embd))
        self.drop = nn.Dropout(config.embd_pdrop)
        mask = torch.tril(torch.ones(config.block_size, config.block_size)).view(
            1, 1, config.block_size, config.block_size
        )
        self.blocks = nn.ModuleList(
            [LayoutActionBlock(config, mask) for _ in range(config.n_layer)]
        )
        self.ln_f = nn.LayerNorm(config.n_embd)
        self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self.block_size = config.block_size
        self.all_tied_weights_keys = dict(self._tied_weights_keys)
        self.post_init()

    def get_block_size(self) -> int:
        """Return the maximum context length."""
        return self.block_size

    def get_input_embeddings(self) -> nn.Embedding:
        """Return token embeddings."""
        return self.tok_emb

    def set_input_embeddings(self, value: nn.Embedding) -> None:
        """Replace token embeddings."""
        self.tok_emb = value

    def forward(
        self,
        input_ids: Int[torch.Tensor, "batch sequence"],
        attention_mask: Bool[torch.Tensor, "batch sequence"] | None = None,
        labels: Int[torch.Tensor, "batch sequence"] | None = None,
        return_dict: bool | None = None,
        output_hidden_states: bool | None = None,
        output_attentions: bool | None = None,
    ) -> CausalLMOutputWithCrossAttentions | tuple[Shaped[torch.Tensor, "..."], ...]:
        """Run a standard causal language-model forward pass.

        Args:
            input_ids: Token ids shaped ``(batch, sequence)``.
            attention_mask: Accepted for Transformers compatibility; causal
                masking follows the checkpoint implementation.
            labels: Optional next-token labels.
            return_dict: Whether to return a dataclass output.
            output_hidden_states: Include final hidden states.
            output_attentions: Accepted for API compatibility; attentions are
                not materialized by the checkpoint-compatible blocks.

        Returns:
            Causal LM output or tuple.

        Raises:
            ValueError: If sequence length exceeds ``block_size``.
        """
        _ = (attention_mask, output_attentions)
        use_return_dict = (
            self.config.use_return_dict if return_dict is None else return_dict
        )
        batch, steps = input_ids.size()
        if steps > self.block_size:
            raise ValueError("Cannot forward; model block size is exhausted.")

        token_embeddings = self.tok_emb(input_ids)
        position_embeddings = self.pos_emb[:, :steps, :]
        hidden_states = self.drop(token_embeddings + position_embeddings)
        for block in self.blocks:
            hidden_states = block(hidden_states)
        hidden_states = self.ln_f(hidden_states)
        logits = self.head(hidden_states)
        loss = None
        if labels is not None:
            target = labels.masked_fill(labels == self.config.pad_token_id, -100)
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                target.view(-1),
                ignore_index=-100,
            )
        if not use_return_dict:
            values: tuple[Shaped[torch.Tensor, "..."], ...]
            values = (logits,) if loss is None else (loss, logits)
            if output_hidden_states:
                values = (*values, hidden_states)
            return values
        return CausalLMOutputWithCrossAttentions(
            loss=cast(torch.FloatTensor | None, loss),
            logits=logits,
            hidden_states=(hidden_states,) if output_hidden_states else None,
        )

    @torch.no_grad()
    def generate(
        self,
        input_ids: Int[torch.Tensor, "batch sequence"],
        *,
        max_new_tokens: int,
        temperature: float = 1.0,
        top_k: int | None = None,
        do_sample: bool = False,
        forced_token_ids: Int[torch.Tensor, "batch sequence"] | None = None,
        generator: torch.Generator | None = None,
    ) -> Int[torch.Tensor, "batch sequence"]:
        """Generate token ids with the reference sampling loop.

        Args:
            input_ids: Prompt token ids.
            max_new_tokens: Number of new tokens.
            temperature: Sampling temperature.
            top_k: Optional top-k crop size.
            do_sample: Whether to use multinomial sampling. If ``False``, greedy
                decoding is used.
            forced_token_ids: Optional ids to force at each generation step.
            generator: Optional torch generator for multinomial sampling.

        Returns:
            Prompt plus generated token ids.
        """
        mode = (
            LayoutActionSamplingMode.top_k
            if do_sample and top_k is not None
            else LayoutActionSamplingMode.multinomial
            if do_sample
            else LayoutActionSamplingMode.greedy
        )
        return sample_action_tokens(
            self,
            input_ids,
            max_new_tokens=max_new_tokens,
            sampling=LayoutActionSamplingConfig(
                mode=mode,
                temperature=temperature,
                top_k=top_k,
            ),
            forced_token_ids=forced_token_ids,
            generator=generator,
        )

class LayoutActionPipeline(LayoutGenerationPipeline):
    """Compose a LayoutAction model and processor for layout generation.

    Args:
        model: Converted LayoutAction causal LM.
        processor: Matching processor/tokenizer.
        config: Optional root pipeline config. Defaults to ``model.config``.

    Examples:
        >>> config = LayoutActionConfig(n_layer=1, n_head=2, n_embd=16, max_elements=1)
        >>> pipe = LayoutActionPipeline(
        ...     model=LayoutActionForCausalLM(config),
        ...     processor=LayoutActionProcessor(LayoutActionTokenizer(config)),
        ...     config=config,
        ... )
        >>> pipe.config.model_type
        'layout-action'
    """

    config_class: ClassVar[type[PretrainedConfig]] = LayoutActionConfig
    component_specs: ClassVar[dict[str, PipelineComponentSpec]] = (
        model_processor_component_specs(
            model_loader=_load_model_component,
            processor_loader=_load_processor_component,
        )
    )

    config: LayoutActionConfig
    model: LayoutActionForCausalLM
    processor: LayoutActionProcessor

    def __init__(
        self,
        model: LayoutActionForCausalLM,
        processor: LayoutActionProcessor,
        config: LayoutActionConfig | None = None,
    ) -> None:
        """Initialize the pipeline."""
        super().__init__(config or model.config)
        self.config = config or model.config
        self.model = model
        self.processor = processor

    @classmethod
    def _from_pretrained_components(
        cls,
        *,
        config: PretrainedConfig,
        components: Mapping[str, PipelineComponent | None],
    ) -> LayoutActionPipeline:
        """Build a pipeline from loaded components."""
        return cls(
            config=cast(LayoutActionConfig, config),
            model=cast(LayoutActionForCausalLM, components["model"]),
            processor=cast(LayoutActionProcessor, components["processor"]),
        )

    @torch.no_grad()
    def __call__(
        self,
        *,
        batch_size: int = 1,
        seed: int | None = None,
        generator: torch.Generator | None = None,
        condition_type: ConditionType | str = ConditionType.unconditional,
        labels: Int[torch.Tensor, "batch elements"]
        | Int[np.ndarray, "batch elements"]
        | Sequence[ArrayLikeInput]
        | None = None,
        bbox: Float[torch.Tensor, "batch elements 4"]
        | Float[np.ndarray, "batch elements 4"]
        | Sequence[ArrayLikeInput]
        | None = None,
        mask: Bool[torch.Tensor, "batch elements"]
        | Bool[np.ndarray, "batch elements"]
        | Sequence[ArrayLikeInput]
        | None = None,
        num_elements: int | list[int] | Int[torch.Tensor, "batch"] | None = None,
        box_format: BoxFormat | str = BoxFormat.xywh,
        normalized: bool = True,
        canvas_size: tuple[int, int] | None = None,
        num_inference_steps: int | None = None,
        output_type: OutputType = "dataclass",
        return_intermediates: bool = False,
        sampling: Literal["greedy", "multinomial", "top_k"] = "top_k",
        temperature: float = 1.0,
        top_k: int | None = 5,
    ) -> LayoutGenerationOutput | LayoutActionOutputDict:  # ty: ignore[invalid-method-override]
        """Generate a layout through the public LayoutAction interface."""
        encoded = self.processor(
            condition_type=condition_type,
            bbox=bbox,
            labels=labels,
            mask=mask,
            num_elements=num_elements,
            box_format=box_format,
            normalized=normalized,
            canvas_size=canvas_size,
            batch_size=batch_size,
            return_tensors="pt",
        )
        model_device = next(self.model.parameters()).device
        prepared_generator = self.prepare_generator(
            generator=generator,
            seed=seed,
            device=model_device,
        )
        input_ids = encoded["input_ids"].to(model_device)
        forced_token_ids = encoded.get("forced_token_ids")
        if isinstance(forced_token_ids, torch.Tensor):
            forced_token_ids = forced_token_ids.to(model_device)
        max_new_tokens = (
            int(num_inference_steps)
            if num_inference_steps is not None
            else int(encoded["max_new_tokens"])
        )
        was_training = self.model.training
        self.model.eval()
        try:
            sequences = self.model.generate(
                input_ids,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_k=top_k,
                do_sample=sampling != "greedy",
                forced_token_ids=forced_token_ids,
                generator=prepared_generator,
            )
        finally:
            self.model.train(was_training)
        return self.processor.post_process_layouts(
            sequences,
            output_type=output_type,
            return_intermediates=return_intermediates,
        )

class LayoutActionProcessor(ProcessorMixin):
    """Prepare LayoutAction prompts and decode generated action sequences.

    Args:
        tokenizer: LayoutAction tokenizer.

    Examples:
        >>> processor = LayoutActionProcessor(LayoutActionTokenizer(LayoutActionConfig(max_elements=1)))
        >>> encoded = processor(condition_type="unconditional")
        >>> encoded["input_ids"].shape
        torch.Size([1, 1])
    """

    attributes = ["tokenizer"]
    tokenizer_class = "LayoutActionTokenizer"

    def __init__(self, tokenizer: LayoutActionTokenizer) -> None:
        """Initialize the processor."""
        self.tokenizer = tokenizer
        self.chat_template = None

    @property
    def config(self) -> LayoutActionConfig:
        """Return the paired LayoutAction config."""
        return self.tokenizer.config

    def __call__(
        self,
        *,
        condition_type: ConditionType | str = ConditionType.unconditional,
        bbox: Float[torch.Tensor, "batch elements 4"]
        | Float[np.ndarray, "batch elements 4"]
        | Sequence[ArrayLikeInput]
        | None = None,
        labels: Int[torch.Tensor, "batch elements"]
        | Int[np.ndarray, "batch elements"]
        | Sequence[ArrayLikeInput]
        | None = None,
        mask: Bool[torch.Tensor, "batch elements"]
        | Bool[np.ndarray, "batch elements"]
        | Sequence[ArrayLikeInput]
        | None = None,
        num_elements: int | list[int] | Int[torch.Tensor, "batch"] | None = None,
        box_format: BoxFormat | str = BoxFormat.xywh,
        normalized: bool = True,
        canvas_size: tuple[int, int] | None = None,
        batch_size: int = 1,
        return_tensors: Literal["pt"] = "pt",
    ) -> BatchEncoding:
        """Encode a public generation condition.

        Args:
            condition_type: Canonical condition or supported release alias.
            bbox: Optional public boxes for completion prompts.
            labels: Optional labels for label/completion prompts.
            mask: Optional valid-element mask.
            num_elements: Optional element count or completion prefix length.
            box_format: Input box format.
            normalized: Whether boxes are normalized.
            canvas_size: Canvas size for pixel boxes.
            batch_size: Batch size for unconditional generation.
            return_tensors: Tensor framework. Only ``pt`` is supported.

        Returns:
            Batch encoding with prompt ids and optional forced token ids.

        Raises:
            NotImplementedError: If the condition is unsupported by LayoutAction.
            ValueError: If required condition payloads are missing.
        """
        if return_tensors != "pt":
            raise ValueError("LayoutActionProcessor only supports return_tensors='pt'")

        condition = normalize_condition_type(condition_type)
        if condition not in SUPPORTED_CONDITIONS:
            raise NotImplementedError(f"LayoutAction does not support {condition}.")

        if condition is ConditionType.unconditional:
            input_ids = torch.full(
                (int(batch_size), 1), self.config.bos_token_id, dtype=torch.long
            )
            return BatchEncoding(
                {
                    "input_ids": input_ids,
                    "attention_mask": torch.ones_like(input_ids),
                    "max_new_tokens": self.config.max_token_length,
                }
            )
        if labels is None:
            raise ValueError(f"{condition} generation requires labels")

        if bbox is None:
            label_tensor = torch.as_tensor(labels)
            if label_tensor.ndim == 1:
                label_tensor = label_tensor.unsqueeze(0)
            mask_tensor = (
                torch.ones_like(label_tensor, dtype=torch.bool)
                if mask is None
                else torch.as_tensor(mask, dtype=torch.bool)
            )
            if mask_tensor.ndim == 1:
                mask_tensor = mask_tensor.unsqueeze(0)
            bbox_tensor = torch.zeros(
                (*label_tensor.shape, 4),
                dtype=torch.float32,
                device=label_tensor.device,
            )
        else:
            bbox_tensor, label_tensor, mask_tensor = prepare_layout_tensors(
                bbox=bbox,
                labels=self._labels_to_ids(labels),
                mask=mask,
                box_format=box_format,
                normalized=normalized,
                canvas_size=canvas_size,
                clamp_converted_normalized=True,
            )
        full = self.tokenizer.encode_layout(
            bbox=bbox_tensor,
            labels=label_tensor.long(),
            mask=mask_tensor.bool(),
        )
        if condition is ConditionType.label:
            input_ids = full[:, :1]
            forced = torch.full(
                (full.size(0), self.config.max_token_length),
                -100,
                dtype=torch.long,
                device=full.device,
            )
            for step in range(
                0,
                self.config.max_elements * self.config.element_token_width,
                self.config.element_token_width,
            ):
                source_index = step + 1
                if source_index < full.size(1):
                    forced[:, step] = full[:, source_index]
            return BatchEncoding(
                {
                    "input_ids": input_ids,
                    "attention_mask": torch.ones_like(input_ids),
                    "forced_token_ids": forced,
                    "max_new_tokens": self.config.max_token_length,
                }
            )
        prefix_elements = self._prefix_elements(num_elements, mask_tensor)
        prefix_length = 1 + prefix_elements * self.config.element_token_width
        input_ids = full[:, :prefix_length]
        remaining = max(0, self.config.max_token_length + 1 - input_ids.size(1))
        return BatchEncoding(
            {
                "input_ids": input_ids,
                "attention_mask": torch.ones_like(input_ids),
                "max_new_tokens": remaining,
            }
        )

    def post_process_layouts(
        self,
        sequences: Int[torch.Tensor, "batch tokens"],
        *,
        output_type: OutputType = "dataclass",
        return_intermediates: bool = False,
    ) -> LayoutGenerationOutput | LayoutActionOutputDict:
        """Decode generated sequences to the common output schema."""
        decoded = self.tokenizer.decode_action_tokens(
            sequences.detach().cpu(),
            return_actions=return_intermediates,
        )
        intermediates = None
        if return_intermediates:
            intermediates = {"actions": decoded.get("actions")}
        output = LayoutGenerationOutput(
            bbox=cast(Float[torch.Tensor, "batch elements 4"], decoded["bbox"]),
            labels=cast(Int[torch.Tensor, "batch elements"], decoded["labels"]),
            mask=cast(Bool[torch.Tensor, "batch elements"], decoded["mask"]),
            id2label=dict(cast(dict[int, str], self.config.id2label)),
            sequences=sequences.detach().cpu(),
            intermediates=intermediates,
        )
        if output_type == "dict":
            return cast(LayoutActionOutputDict, dict(output))
        if output_type == "dataclass":
            return output
        raise ValueError(f"Unsupported output_type: {output_type}")

    def save_pretrained(
        self,
        save_directory: str | PathLike[str],
        push_to_hub: bool = False,
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Save processor and tokenizer metadata."""
        _ = (push_to_hub, kwargs)
        out_dir = Path(save_directory)
        out_dir.mkdir(parents=True, exist_ok=True)
        self.tokenizer.save_pretrained(out_dir)
        with (out_dir / PROCESSOR_CONFIG_FILE).open("w", encoding="utf-8") as f:
            json.dump({"processor_class": self.__class__.__name__}, f, indent=2)

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: str | PathLike[str],
        cache_dir: str | PathLike[str] | None = None,
        force_download: bool = False,
        local_files_only: bool = False,
        token: str | bool | None = None,
        revision: str = "main",
        **kwargs: str | int | float | bool | None,
    ) -> "LayoutActionProcessor":
        """Load processor metadata from a checkpoint directory or Hub repo id."""
        tokenizer = LayoutActionTokenizer.from_pretrained(
            pretrained_model_name_or_path,
            cache_dir=cache_dir,
            force_download=force_download,
            local_files_only=local_files_only,
            token=token,
            revision=revision,
            **kwargs,
        )
        return cls(tokenizer=tokenizer)

    def _labels_to_ids(
        self,
        labels: Int[torch.Tensor, "batch elements"]
        | Int[np.ndarray, "batch elements"]
        | Sequence[ArrayLikeInput],
    ) -> Int[torch.Tensor, "batch elements"]:
        if isinstance(labels, torch.Tensor):
            return labels.long()
        label_array = np.asarray(labels, dtype=object)
        if all(not isinstance(label, str) for label in label_array.flatten()):
            return torch.as_tensor(labels, dtype=torch.long)
        label2id = cast(dict[str, int], self.config.label2id)
        vectorized = np.vectorize(lambda label: label2id[str(label)])
        return torch.as_tensor(vectorized(label_array), dtype=torch.long)

    def _prefix_elements(
        self,
        num_elements: int | list[int] | Int[torch.Tensor, "batch"] | None,
        mask: Bool[torch.Tensor, "batch elements"],
    ) -> int:
        if num_elements is None:
            valid_counts = mask.sum(dim=1)
            return int(valid_counts.min().item())
        if isinstance(num_elements, int):
            return int(num_elements)
        tensor = torch.as_tensor(num_elements)
        return int(tensor.min().item())