122 lines
5.3 KiB
Python
122 lines
5.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""MiniMax H3 text-conditioning contract."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Mapping
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import torch
|
|
|
|
MINIMAX_H3_TEXT_CONDITIONING_SCHEMA = "minimax_h3.text_conditioning/v1"
|
|
MINIMAX_H3_TEXT_HIDDEN_SIZE = 5120
|
|
MINIMAX_H3_PRESENTATION_TASK_KEY = "_minimax_h3_presentation_task"
|
|
MINIMAX_H3_CONDITION_LABELS_KEY = "_minimax_h3_condition_labels"
|
|
|
|
|
|
def _validate_contiguous_strided_layout(name: str, tensor: torch.Tensor) -> None:
|
|
if tensor.layout != torch.strided:
|
|
raise ValueError(
|
|
f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: {name} must use contiguous strided layout, got {tensor.layout}"
|
|
)
|
|
if not tensor.is_contiguous():
|
|
raise ValueError(
|
|
f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: {name} must use contiguous "
|
|
f"strided layout, got stride={tuple(tensor.stride())}"
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class MiniMaxH3TextConditioning:
|
|
"""``minimax_h3.text_conditioning/v1`` semantic payload."""
|
|
|
|
hidden_states: torch.Tensor
|
|
token_tags: torch.Tensor
|
|
|
|
@classmethod
|
|
def from_payload(
|
|
cls,
|
|
payload: Mapping[str, Any],
|
|
) -> MiniMaxH3TextConditioning:
|
|
"""Validate the semantic payload consumed by the diffusion stage."""
|
|
hidden_states = payload.get("hidden_states")
|
|
token_tags = payload.get("token_tags")
|
|
if not isinstance(hidden_states, torch.Tensor) or not isinstance(token_tags, torch.Tensor):
|
|
raise ValueError(
|
|
f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: conditioning requires hidden_states and token_tags tensors"
|
|
)
|
|
if hidden_states.ndim != 2 or hidden_states.shape[-1] != MINIMAX_H3_TEXT_HIDDEN_SIZE:
|
|
raise ValueError(
|
|
f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: hidden_states must have shape "
|
|
f"[tokens, {MINIMAX_H3_TEXT_HIDDEN_SIZE}], got {tuple(hidden_states.shape)}"
|
|
)
|
|
if hidden_states.dtype != torch.bfloat16:
|
|
raise ValueError(
|
|
f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: hidden_states must have dtype "
|
|
f"torch.bfloat16, got {hidden_states.dtype}"
|
|
)
|
|
_validate_contiguous_strided_layout("hidden_states", hidden_states)
|
|
if token_tags.ndim != 1 or token_tags.shape[0] != hidden_states.shape[0]:
|
|
raise ValueError(
|
|
f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: token_tags must align with hidden_states, got "
|
|
f"token_tags={tuple(token_tags.shape)} and hidden_states={tuple(hidden_states.shape)}"
|
|
)
|
|
if token_tags.dtype != torch.int64:
|
|
raise ValueError(
|
|
f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: token_tags must have dtype torch.int64, got {token_tags.dtype}"
|
|
)
|
|
_validate_contiguous_strided_layout("token_tags", token_tags)
|
|
if not torch.all((token_tags == 0) | (token_tags == 1)):
|
|
raise ValueError(
|
|
f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: text-encoder token_tags must contain only 0 and 1"
|
|
)
|
|
return cls(hidden_states=hidden_states, token_tags=token_tags)
|
|
|
|
@classmethod
|
|
def from_omni_payload(
|
|
cls,
|
|
payload: Mapping[str, Any],
|
|
) -> MiniMaxH3TextConditioning:
|
|
"""Validate and adapt the existing ``OmniPayload`` stage-wire view."""
|
|
hidden_states = payload.get("hidden_states")
|
|
if not isinstance(hidden_states, Mapping):
|
|
raise ValueError(f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: text encoder returned no hidden_states payload")
|
|
hidden = hidden_states.get("output")
|
|
if not isinstance(hidden, torch.Tensor):
|
|
raise ValueError(
|
|
f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: text encoder returned no hidden_states.output tensor"
|
|
)
|
|
|
|
meta = payload.get("meta")
|
|
if not isinstance(meta, Mapping):
|
|
raise ValueError(f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: text encoder returned no conditioning metadata")
|
|
token_role_ids = meta.get("token_role_ids")
|
|
if not isinstance(token_role_ids, torch.Tensor):
|
|
raise ValueError(f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: text encoder returned no token_role_ids tensor")
|
|
if token_role_ids.ndim != 2 or token_role_ids.shape[-1] != 1:
|
|
raise ValueError(
|
|
f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: stage-wire token_role_ids must have shape "
|
|
f"[tokens, 1], got {tuple(token_role_ids.shape)}"
|
|
)
|
|
if token_role_ids.dtype != torch.int64:
|
|
raise ValueError(
|
|
f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: stage-wire token_role_ids must have dtype "
|
|
f"torch.int64, got {token_role_ids.dtype}"
|
|
)
|
|
_validate_contiguous_strided_layout("stage-wire token_role_ids", token_role_ids)
|
|
|
|
return cls.from_payload(
|
|
{
|
|
"hidden_states": hidden,
|
|
"token_tags": token_role_ids.squeeze(-1),
|
|
}
|
|
)
|
|
|
|
def to_payload(self) -> dict[str, torch.Tensor]:
|
|
return {
|
|
"hidden_states": self.hidden_states,
|
|
"token_tags": self.token_tags,
|
|
}
|