timebraid/model/timebraid.py

"""HF-native TimeBraid model with directly registered language and TS components."""

from __future__ import annotations

from collections.abc import Iterable, Mapping, Sequence
from dataclasses import dataclass, fields
from typing import TYPE_CHECKING, Any, Optional, Union

import torch
import torch.nn as nn
from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    PretrainedConfig,
    PreTrainedModel,
)
from transformers.cache_utils import Cache
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.modeling_utils import get_parameter_dtype
from transformers.utils import ModelOutput

from .mot.bridge_runtime import initialize_timebraid_components
from .mot.bridge_runtime_mot import sync_gradient_checkpointing_from_decoder_layers
from .mot.config_contract import (
    MOT_HF_CONFIG_CONTRACT_FIELDS,
    collect_mot_hf_config_contract,
)
from .mot.generation_route import (
    GenerationRoute,
    require_mixed_generation_inputs,
    resolve_generation_route,
)
from .mot.kv_cache import MoTDynamicCache
from .mot.model import (
    NormalizedMoTGenerationControls,
    TimeBraidGenerateOutput,
    TimeBraidPayload,
    merge_timebraid_losses,
    normalize_mot_generation_controls,
    run_timebraid_decoder,
    run_timebraid_generate,
    validate_mot_generation_input_ids,
)
from .mot.structures import IGNORE_INDEX, TIMEBRAID_TS_PAYLOAD_FIELDS

if TYPE_CHECKING:  # pragma: no cover - import cycle only matters at runtime
    from ..processing_timebraid import TimeBraidBatchFeature

_FINITE_SCAN_CHUNK_ELEMENTS = 8 * 1024 * 1024
_NON_INFERENCE_CONFIG_FIELDS = (
    "mot_timesfm_model_name_or_path",
    "mot_ts_roi_mse_alpha",
    "mot_ts_loss_weight",
    "mot_lm_loss_weight",
    "mot_ts_understanding_loss_weight",
)
_FORBIDDEN_CONFIG_FIELDS = _NON_INFERENCE_CONFIG_FIELDS
_MOT_GENERATION_CONTROL_FIELDS = (
    "mot_target_horizons",
    "mot_forecast_head_len",
    "mot_target_total_lengths",
    "mot_target_history_span_idxs",
    "mot_return_forecast_quantiles",
)


@dataclass
class TimeBraidLossBreakdown(ModelOutput):
    """Fixed scalar loss accounting returned beside the differentiable loss."""

    loss: Optional[torch.Tensor] = None
    lm_loss: Optional[torch.Tensor] = None
    ts_aux_loss: Optional[torch.Tensor] = None
    ts_aux_point_loss: Optional[torch.Tensor] = None
    ts_aux_quantile_loss: Optional[torch.Tensor] = None
    ts_aux_global_point_loss: Optional[torch.Tensor] = None
    ts_aux_global_quantile_loss: Optional[torch.Tensor] = None
    ts_aux_global_uncapped_loss: Optional[torch.Tensor] = None
    ts_aux_global_cap_scale: Optional[torch.Tensor] = None
    ts_aux_roi_loss: Optional[torch.Tensor] = None
    ts_aux_roi_point_loss: Optional[torch.Tensor] = None
    ts_aux_roi_weighted_point_loss: Optional[torch.Tensor] = None
    ts_aux_roi_uncapped_loss: Optional[torch.Tensor] = None
    ts_aux_roi_weighted_uncapped_loss: Optional[torch.Tensor] = None
    ts_aux_roi_cap_scale: Optional[torch.Tensor] = None
    ts_aux_roi_supervised_points: Optional[torch.Tensor] = None
    ts_aux_roi_supervised_targets: Optional[torch.Tensor] = None
    ts_aux_roi_alpha: Optional[torch.Tensor] = None
    ts_aux_understanding_loss: Optional[torch.Tensor] = None
    ts_aux_understanding_weighted_loss: Optional[torch.Tensor] = None
    ts_aux_understanding_time_loss: Optional[torch.Tensor] = None
    ts_aux_understanding_fft_loss: Optional[torch.Tensor] = None
    ts_aux_understanding_patches: Optional[torch.Tensor] = None
    ts_aux_understanding_points: Optional[torch.Tensor] = None
    ts_aux_understanding_weight: Optional[torch.Tensor] = None


@dataclass
class TimeBraidOutput(CausalLMOutputWithPast):
    """Stable causal-LM output for every dictionary-returning TimeBraid forward."""

    token_accuracy: Optional[torch.Tensor] = None
    predicted_tokens: Optional[torch.Tensor] = None
    loss_breakdown: Optional[TimeBraidLossBreakdown] = None


def _present_forbidden_config_fields(config: object) -> list[str]:
    if isinstance(config, Mapping):
        return sorted(field for field in _FORBIDDEN_CONFIG_FIELDS if field in config)
    return sorted(field for field in _FORBIDDEN_CONFIG_FIELDS if hasattr(config, field))


def _validate_finite_named_tensors(
    named_tensors: Iterable[tuple[str, torch.Tensor]],
    *,
    owner: str,
) -> None:
    for name, tensor in named_tensors:
        if not tensor.is_floating_point() and not tensor.is_complex():
            continue
        flat = tensor.detach().contiguous().view(-1)
        for start in range(0, int(flat.numel()), _FINITE_SCAN_CHUNK_ELEMENTS):
            chunk = flat[start : start + _FINITE_SCAN_CHUNK_ELEMENTS]
            finite_mask = torch.isfinite(chunk)
            if bool(torch.all(finite_mask)):
                continue
            chunk_index = int(
                torch.nonzero(torch.logical_not(finite_mask), as_tuple=False)[0]
                .detach()
                .cpu()
                .item()
            )
            raise RuntimeError(
                f"{owner} contains a non-finite tensor value: "
                f"name={name!r}, flat_index={start + chunk_index}, dtype={tensor.dtype}."
            )


def _validate_finite_model_tensors(model: nn.Module, *, owner: str) -> None:
    """Reject non-finite floating or complex model state before use."""
    _validate_finite_named_tensors(
        (*model.named_parameters(), *model.named_buffers()),
        owner=owner,
    )


def _runtime_options_from_timebraid_config(config: "TimeBraidConfig") -> dict[str, Any]:
    """Build direct-component initialization options from checkpoint metadata.

    A contract field's name *is* its runtime option name, so this copies
    rather than translates. The hand-maintained rename table it replaced was
    what allowed one knob to carry two spellings and drift apart.
    """
    contract = collect_mot_hf_config_contract(config)
    options: dict[str, Any] = {
        name: contract[name] for name in MOT_HF_CONFIG_CONTRACT_FIELDS
    }
    # Never restored from a checkpoint: inference always runs with gradient
    # checkpointing off.
    options["mot_ts_gradient_checkpointing"] = False
    options["mot_ts_gradient_checkpointing_use_reentrant"] = False
    compute_dtype = getattr(config, "dtype", None) or getattr(
        config.llm_config, "dtype", None
    )
    if compute_dtype is not None:
        options["mot_compute_dtype"] = compute_dtype
    return options


_LLM_MIRROR_FIELDS = (
    "vocab_size",
    "hidden_size",
    "intermediate_size",
    "num_hidden_layers",
    "num_attention_heads",
    "num_key_value_heads",
    "head_dim",
    "max_position_embeddings",
    "tie_word_embeddings",
    "dtype",
    "bos_token_id",
    "eos_token_id",
    "pad_token_id",
    "use_cache",
    "sliding_window",
    "layer_types",
    "_attn_implementation",
    "_attn_implementation_internal",
)
_LLM_SERIALIZED_MIRROR_FIELDS = tuple(
    field for field in _LLM_MIRROR_FIELDS if not field.startswith("_")
)


class TimeBraidConfig(PretrainedConfig):
    r"""Composite config for TimeBraid's language and time-series components."""

    model_type = "timebraid"
    is_composition = True
    has_no_defaults_at_init = True
    sub_configs = {"llm_config": AutoConfig}

    def __init__(
        self,
        llm_config: Optional[PretrainedConfig | dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        if llm_config is None:
            raise ValueError("TimeBraidConfig requires `llm_config`.")
        forbidden_fields = _present_forbidden_config_fields(kwargs)
        if forbidden_fields:
            raise ValueError(
                "TimeBraidConfig does not restore training-only or external "
                f"TimesFM fields: {forbidden_fields}."
            )
        nested_forbidden_fields = _present_forbidden_config_fields(llm_config)
        if nested_forbidden_fields:
            raise ValueError(
                "TimeBraidConfig llm_config does not restore training-only or "
                f"external TimesFM fields: {nested_forbidden_fields}."
            )

        if isinstance(llm_config, PretrainedConfig):
            resolved_llm_config = llm_config
        else:
            llm_config_dict = dict(llm_config)
            llm_model_type = llm_config_dict.pop("model_type", None)
            if llm_model_type is None:
                raise ValueError(
                    "TimeBraidConfig `llm_config` must include `model_type`."
                )
            resolved_llm_config = AutoConfig.for_model(
                llm_model_type, **llm_config_dict
            )

        outer_mirrors = sorted(
            field for field in _LLM_SERIALIZED_MIRROR_FIELDS if field in kwargs
        )
        if outer_mirrors:
            raise ValueError(
                "TimeBraidConfig stores LLM fields only in llm_config; remove outer fields "
                f"{outer_mirrors}."
            )
        kwargs.setdefault("architectures", ["TimeBraid"])
        # PretrainedConfig applies dtype and attention overrides during its
        # initializer. Register the nested config first so those composite
        # controls reach Qwen instead of remaining only on the outer wrapper.
        self.llm_config = resolved_llm_config
        super().__init__(**kwargs)
        self._sync_llm_surface()

    @classmethod
    def from_llm_config(
        cls, llm_config: PretrainedConfig, **kwargs: Any
    ) -> "TimeBraidConfig":
        return cls(llm_config=llm_config, **kwargs)

    def _sync_llm_surface(self) -> None:
        """Derive the non-serialized wrapper surface generic HF helpers consume."""
        for field in _LLM_MIRROR_FIELDS:
            if hasattr(self.llm_config, field):
                setattr(self, field, getattr(self.llm_config, field))

    def to_dict(self) -> dict[str, Any]:
        """Serialize only inference restore state, never training objectives."""
        payload = super().to_dict()
        forbidden_fields = _present_forbidden_config_fields(payload)
        if forbidden_fields:
            raise RuntimeError(
                "TimeBraidConfig contains forbidden training-only or external "
                f"TimesFM fields: {forbidden_fields}."
            )
        for field in _LLM_SERIALIZED_MIRROR_FIELDS:
            payload.pop(field, None)
        nested = payload.get("llm_config")
        if isinstance(nested, dict):
            nested_forbidden_fields = _present_forbidden_config_fields(nested)
            if nested_forbidden_fields:
                raise RuntimeError(
                    "TimeBraidConfig llm_config contains forbidden training-only, "
                    f"or external TimesFM fields: {nested_forbidden_fields}."
                )
        return payload


def _labels_with_language_supervision(
    labels: Optional[torch.LongTensor],
) -> Optional[torch.LongTensor]:
    if isinstance(labels, torch.Tensor) and (
        labels.numel() == 0 or not bool(torch.any(labels.ne(IGNORE_INDEX)).item())
    ):
        return None
    return labels


def _require_binary_plain_attention_mask(attention_mask: object) -> None:
    """Reject packed segment ids before delegating a plain call to native Qwen."""
    # Rank-4 prepared masks remain part of Qwen's public ABI. TimeBraid packing
    # uses positive segment ids only in the rank-2 mask produced by its
    # processor/collator, so the plain route must accept exactly binary values
    # there instead of letting Qwen collapse every nonzero id to one segment.
    if not isinstance(attention_mask, torch.Tensor) or attention_mask.ndim != 2:
        return
    binary_values = torch.logical_or(attention_mask.eq(0), attention_mask.eq(1))
    if not bool(torch.all(binary_values).item()):
        raise RuntimeError(
            "Plain Qwen fallback requires a binary 2D attention_mask with values "
            "in {0, 1}; packed segment ids require the TimeBraid processor/collator "
            f"payload, got dtype={attention_mask.dtype}, shape={tuple(attention_mask.shape)}."
        )


_TIMEBRAID_KWARG_PREFIXES = ("mot_", "ts_")
_KNOWN_TIMEBRAID_KWARGS = frozenset(
    TIMEBRAID_TS_PAYLOAD_FIELDS + _MOT_GENERATION_CONTROL_FIELDS
)


def _reject_unknown_timebraid_kwargs(generation_kwargs: Mapping[str, object]) -> None:
    """Reject misspelled TimeBraid arguments instead of forwarding them.

    `mot_` and `ts_` are TimeBraid's own argument namespace, so an unrecognized
    name there is a mistake we can name. Unprefixed names belong to Hugging
    Face and must keep flowing through, or this package would break whenever
    transformers adds a generation argument.

    Without this, a typo such as `mot_target_horizon` was silently forwarded,
    the request routed to plain text, and the caller received prose with no
    forecast and no error.
    """
    unknown = sorted(
        name
        for name in generation_kwargs
        if name.startswith(_TIMEBRAID_KWARG_PREFIXES)
        and name not in _KNOWN_TIMEBRAID_KWARGS
    )
    if unknown:
        raise TypeError(
            f"Unknown TimeBraid generation argument(s) {unknown}. Supported "
            f"names are {sorted(_KNOWN_TIMEBRAID_KWARGS)}."
        )


def _fixed_loss_breakdown(
    breakdown: Optional[Mapping[str, torch.Tensor]],
) -> Optional[TimeBraidLossBreakdown]:
    if breakdown is None:
        return None
    reference = next(
        (value for value in breakdown.values() if isinstance(value, torch.Tensor)),
        None,
    )
    if reference is None:
        raise RuntimeError("TimeBraid loss accounting returned no tensor values.")
    zero = reference.detach().new_zeros(())
    values: dict[str, torch.Tensor] = {}
    for output_field in fields(TimeBraidLossBreakdown):
        value = breakdown.get(output_field.name)
        if value is None:
            value = zero.clone()
        if not isinstance(value, torch.Tensor) or value.numel() != 1:
            raise RuntimeError(
                "TimeBraid loss accounting values must be scalar tensors: "
                f"field={output_field.name!r}, value={value!r}."
            )
        values[output_field.name] = value.detach().reshape(())
    return TimeBraidLossBreakdown(**values)


def _as_timebraid_output(
    outputs: object,
    breakdown: Optional[Mapping[str, torch.Tensor]],
) -> TimeBraidOutput:
    return TimeBraidOutput(
        loss=getattr(outputs, "loss", None),
        logits=getattr(outputs, "logits", None),
        past_key_values=getattr(outputs, "past_key_values", None),
        hidden_states=getattr(outputs, "hidden_states", None),
        attentions=getattr(outputs, "attentions", None),
        token_accuracy=getattr(outputs, "token_accuracy", None),
        predicted_tokens=getattr(outputs, "predicted_tokens", None),
        loss_breakdown=_fixed_loss_breakdown(breakdown),
    )


def _build_lm_output(
    *,
    llm: PreTrainedModel,
    hidden_states: torch.Tensor,
    past_key_values: object,
    decoder_hidden_states: object,
    decoder_attentions: object,
    labels: Optional[torch.LongTensor],
    labels_for_loss: Optional[torch.LongTensor],
    shift_labels: Optional[torch.LongTensor],
    logits_to_keep: Union[int, torch.Tensor],
    skip_logits: Optional[bool],
    loss_kwargs: Mapping[str, object],
) -> TimeBraidOutput:
    """Project decoder hidden states and own the complete language-loss contract."""
    if skip_logits is not None and type(skip_logits) is not bool:
        raise TypeError(
            f"TimeBraid forward skip_logits must be a bool or None, got {skip_logits!r}."
        )

    # `skip_logits=True` only ever meant "the fused Liger loss will consume the
    # hidden states instead of logits". That path was unreachable — the flag
    # enabling it was hardcoded False with no setter — so it is gone, and the
    # request is refused rather than silently ignored.
    if skip_logits:
        raise RuntimeError(
            "skip_logits=True is not supported: TimeBraid always projects logits "
            "through the language head."
        )

    slice_indices = (
        slice(-logits_to_keep, None)
        if isinstance(logits_to_keep, int)
        else logits_to_keep
    )
    kept_hidden_states = hidden_states[:, slice_indices, :]
    has_effective_supervision = labels_for_loss is not None or shift_labels is not None
    loss = None
    # Only the removed fused-loss path produced these; the language head does
    # not, so they stay None and keep TimeBraidOutput's shape stable.
    token_accuracy = None
    predicted_tokens = None
    logits = llm.lm_head(kept_hidden_states)
    if has_effective_supervision:
        loss = llm.loss_function(
            logits=logits,
            labels=labels_for_loss,
            shift_labels=shift_labels,
            vocab_size=llm.config.vocab_size,
            **loss_kwargs,
        )

    return TimeBraidOutput(
        loss=loss,
        logits=logits,
        past_key_values=past_key_values,
        hidden_states=decoder_hidden_states,
        attentions=decoder_attentions,
        token_accuracy=token_accuracy,
        predicted_tokens=predicted_tokens,
    )


class TimeBraid(PreTrainedModel):
    r"""Top-level model with direct language, time-series, and attention components."""

    # This is the complete trainable/inference component graph. Keeping the
    # ownership list on the public model makes optional topology readable
    # without inspecting the numeric execution helpers.
    llm: PreTrainedModel
    generation_tsfm: nn.Module
    understanding_tsfm: Optional[nn.Module]
    global_residual_attention: nn.ModuleDict
    packed_attention: nn.Module
    understanding_head: Optional[nn.Module]

    config_class = TimeBraidConfig
    base_model_prefix = "llm"
    main_input_name = "input_ids"
    _supports_flash_attn_2 = True
    _supports_sdpa = True
    supports_gradient_checkpointing = True
    _tied_weights_keys = {"llm.lm_head.weight": "llm.model.embed_tokens.weight"}

    def __init__(
        self,
        config: TimeBraidConfig,
        llm: Optional[PreTrainedModel] = None,
        runtime_options: Optional[dict[str, Any]] = None,
    ) -> None:
        if not isinstance(config, TimeBraidConfig):
            raise ValueError(
                f"TimeBraid requires TimeBraidConfig, got {type(config).__name__}."
            )
        super().__init__(config)
        self.llm = (
            llm
            if llm is not None
            else AutoModelForCausalLM.from_config(config.llm_config)
        )
        self.config.llm_config = self.llm.config
        self.config._sync_llm_surface()
        self.generation_config = self.llm.generation_config

        resolved_runtime_options = _runtime_options_from_timebraid_config(config)
        if runtime_options:
            resolved_runtime_options.update(runtime_options)
        initialize_timebraid_components(
            self,
            self.llm.model.config,
            resolved_runtime_options,
        )
        # Optional training integration may enable the fused hidden-to-loss tail
        # after construction; that tail does not replace either Qwen forward method.
        self.generation_tsfm.eval()
        if self.understanding_tsfm is not None:
            self.understanding_tsfm.eval()
        self._sync_hf_loading_contracts()

    def _sync_hf_loading_contracts(self) -> None:
        """Synchronize tied-weight and parallelism metadata without generic post-init."""
        self._tied_weights_keys = {
            "llm.lm_head.weight": "llm.model.embed_tokens.weight"
        }
        self.all_tied_weights_keys = dict(self._tied_weights_keys)
        tied_keys = getattr(self.llm, "all_tied_weights_keys", None)
        if tied_keys:
            self.all_tied_weights_keys.update(
                {f"llm.{key}": f"llm.{value}" for key, value in tied_keys.items()}
            )
        # Transformers expects composite wrappers to expose concrete empty
        # parallelism plans during distributed `from_pretrained()` allocator
        # warmup, so mirror the child plans here.
        for plan_attr in ("_tp_plan", "_ep_plan", "_pp_plan"):
            merged_plan = {}
            child_plan = getattr(self.llm, plan_attr, None)
            if child_plan:
                merged_plan.update(
                    {f"llm.{key}": value for key, value in child_plan.copy().items()}
                )
            setattr(self, plan_attr, merged_plan)

    def train(self, mode: bool = True) -> "TimeBraid":
        super().train(mode)
        # TimesFM towers stay in eval mode during full-parameter training. Module
        # mode controls stochastic behavior independently of `requires_grad`.
        self.generation_tsfm.eval()
        if self.understanding_tsfm is not None:
            self.understanding_tsfm.eval()
        return self

    def gradient_checkpointing_enable(
        self, gradient_checkpointing_kwargs: Optional[dict[str, Any]] = None
    ) -> None:
        super().gradient_checkpointing_enable(
            gradient_checkpointing_kwargs=gradient_checkpointing_kwargs
        )
        sync_gradient_checkpointing_from_decoder_layers(self, self.llm.model.layers)

    def gradient_checkpointing_disable(self) -> None:
        super().gradient_checkpointing_disable()
        sync_gradient_checkpointing_from_decoder_layers(self, self.llm.model.layers)

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        logits_to_keep: Union[int, torch.Tensor] = 0,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        skip_logits: Optional[bool] = None,
        shift_labels: Optional[torch.LongTensor] = None,
        ts_values: object = None,
        ts_lengths: object = None,
        ts_loss_start_idxs: object = None,
        ts_loss_roi_masks: object = None,
        ts_roles: object = None,
        ts_segment_ids: object = None,
        ts_span_mask: object = None,
        ts_text_start_token_idxs: object = None,
        ts_text_end_token_idxs: object = None,
        active_rows: Optional[torch.Tensor] = None,
        num_items_in_batch: Optional[object] = None,
        mot_ts_num_records_by_horizon: Optional[object] = None,
        mot_ts_num_roi_targets_by_horizon: Optional[object] = None,
        mot_ts_understanding_num_points: Optional[object] = None,
        mot_ts_understanding_num_fft_elements: Optional[object] = None,
        **kwargs: Any,
    ) -> TimeBraidOutput | tuple[Any, ...]:
        # Generation-only controls can accompany processor batches, but never
        # participate in the causal forward computation.
        for key in _MOT_GENERATION_CONTROL_FIELDS:
            kwargs.pop(key, None)

        payload = TimeBraidPayload(
            ts_values=ts_values,
            ts_lengths=ts_lengths,
            ts_loss_start_idxs=ts_loss_start_idxs,
            ts_loss_roi_masks=ts_loss_roi_masks,
            ts_roles=ts_roles,
            ts_segment_ids=ts_segment_ids,
            ts_span_mask=ts_span_mask,
            ts_text_start_token_idxs=ts_text_start_token_idxs,
            ts_text_end_token_idxs=ts_text_end_token_idxs,
        )
        use_timebraid_decoder = (
            isinstance(past_key_values, MoTDynamicCache) or payload.has_runtime_inputs()
        )
        if not use_timebraid_decoder and isinstance(input_ids, torch.Tensor):
            use_timebraid_decoder = bool(
                torch.logical_or(
                    input_ids.eq(int(self.ts_open_token_id)),
                    input_ids.eq(int(self.ts_close_token_id)),
                )
                .any()
                .item()
            )
        if active_rows is not None and not isinstance(past_key_values, MoTDynamicCache):
            raise RuntimeError("active_rows requires an existing MoTDynamicCache.")

        if return_dict is None:
            resolved_return_dict = bool(self.llm.config.use_return_dict)
        elif type(return_dict) is bool:
            resolved_return_dict = return_dict
        else:
            raise TypeError(
                f"TimeBraid forward return_dict must be a bool or None, got {return_dict!r}."
            )
        labels_for_loss = _labels_with_language_supervision(labels)
        if not use_timebraid_decoder:
            _require_binary_plain_attention_mask(attention_mask)
            if skip_logits is not None and type(skip_logits) is not bool:
                raise TypeError(
                    "TimeBraid forward skip_logits must be a bool or None, got "
                    f"{skip_logits!r}."
                )
            use_outer_lm_head = shift_labels is not None and labels_for_loss is None
            if skip_logits and not use_outer_lm_head:
                raise ValueError(
                    "skip_logits=True requires fused TimeBraid language loss."
                )
            plain_loss_kwargs = dict(kwargs)
            if num_items_in_batch is not None:
                plain_loss_kwargs["num_items_in_batch"] = num_items_in_batch
            plain_forward_kwargs = dict(plain_loss_kwargs)
            if shift_labels is not None:
                plain_forward_kwargs["shift_labels"] = shift_labels
            if use_outer_lm_head:
                # Run Qwen's ordinary decoder, then let TimeBraid consume its
                # hidden states. This outer tail exists for explicit
                # shift-label-only loss, where Qwen cannot build the labels.
                decoder_outputs = self.llm.model(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    position_ids=position_ids,
                    past_key_values=past_key_values,
                    inputs_embeds=inputs_embeds,
                    use_cache=use_cache,
                    cache_position=cache_position,
                    output_attentions=output_attentions,
                    output_hidden_states=output_hidden_states,
                    return_dict=True,
                    **kwargs,
                )
                outputs = _build_lm_output(
                    llm=self.llm,
                    hidden_states=decoder_outputs.last_hidden_state,
                    past_key_values=decoder_outputs.past_key_values,
                    decoder_hidden_states=decoder_outputs.hidden_states,
                    decoder_attentions=decoder_outputs.attentions,
                    labels=labels,
                    labels_for_loss=labels_for_loss,
                    shift_labels=shift_labels,
                    logits_to_keep=logits_to_keep,
                    skip_logits=skip_logits,
                    loss_kwargs=plain_loss_kwargs,
                )
            else:
                outputs = self.llm(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    position_ids=position_ids,
                    past_key_values=past_key_values,
                    inputs_embeds=inputs_embeds,
                    labels=labels_for_loss,
                    use_cache=use_cache,
                    cache_position=cache_position,
                    logits_to_keep=logits_to_keep,
                    output_attentions=output_attentions,
                    output_hidden_states=output_hidden_states,
                    return_dict=True,
                    **plain_forward_kwargs,
                )
            outputs, breakdown = merge_timebraid_losses(
                model=self,
                mot_runtime=None,
                outputs=outputs,
            )
            timebraid_output = _as_timebraid_output(outputs, breakdown)
            if resolved_return_dict:
                return timebraid_output
            # Preserve Qwen's tuple prefix and ordering. Logits keeps its slot
            # when fused execution leaves it as None, followed by optional decoder
            # outputs; TimeBraid's token statistics append afterward.
            tuple_output = (timebraid_output.logits,)
            tuple_output += tuple(
                value
                for value in (
                    timebraid_output.past_key_values,
                    timebraid_output.hidden_states,
                    timebraid_output.attentions,
                )
                if value is not None
            )
            if timebraid_output.loss is not None:
                tuple_output = (timebraid_output.loss,) + tuple_output
            if timebraid_output.token_accuracy is not None:
                tuple_output += (timebraid_output.token_accuracy,)
            if timebraid_output.predicted_tokens is not None:
                tuple_output += (timebraid_output.predicted_tokens,)
            return tuple_output

        if not resolved_return_dict:
            raise RuntimeError("TS-routed TimeBraid forward requires return_dict=True.")
        cache_implementation = kwargs.pop("cache_implementation", None)
        if use_cache is None:
            configured_use_cache = self.llm.model.config.use_cache
            if type(configured_use_cache) is not bool:
                raise TypeError(
                    "Qwen config use_cache must be a bool, got "
                    f"{configured_use_cache!r}."
                )
            resolved_use_cache = configured_use_cache and not self.llm.model.training
        elif type(use_cache) is bool:
            resolved_use_cache = use_cache
        else:
            raise TypeError(
                f"TimeBraid forward use_cache must be a bool or None, got {use_cache!r}."
            )
        decoder_result = run_timebraid_decoder(
            owner=self,
            text_model=self.llm.model,
            payload=payload,
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=resolved_use_cache,
            cache_position=cache_position,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=True,
            active_rows=active_rows,
            cache_implementation=cache_implementation,
            collect_incremental_kv_cache=resolved_use_cache,
        )

        loss_kwargs = dict(kwargs)
        if num_items_in_batch is not None:
            loss_kwargs["num_items_in_batch"] = num_items_in_batch
        outputs = _build_lm_output(
            llm=self.llm,
            hidden_states=decoder_result.hidden_states,
            past_key_values=decoder_result.past_key_values,
            decoder_hidden_states=None,
            decoder_attentions=None,
            labels=labels,
            labels_for_loss=labels_for_loss,
            shift_labels=shift_labels,
            logits_to_keep=logits_to_keep,
            skip_logits=skip_logits,
            loss_kwargs=loss_kwargs,
        )
        outputs, breakdown = merge_timebraid_losses(
            model=self,
            mot_runtime=decoder_result.mot_runtime,
            outputs=outputs,
            mot_ts_num_records_by_horizon=mot_ts_num_records_by_horizon,
            mot_ts_num_roi_targets_by_horizon=mot_ts_num_roi_targets_by_horizon,
            mot_ts_understanding_num_points=mot_ts_understanding_num_points,
            mot_ts_understanding_num_fft_elements=(
                mot_ts_understanding_num_fft_elements
            ),
        )
        return _as_timebraid_output(outputs, breakdown)

    def save_pretrained(self, *args: Any, **kwargs: Any) -> None:
        # Match Transformers' own dtype authority, but update both composite
        # config surfaces before the base implementation serializes either one.
        if self.config.llm_config is not self.llm.config:
            raise RuntimeError(
                "TimeBraid save requires config.llm_config to be the live LLM config."
            )
        # The checkpoint is self-contained, so machine-local source paths are
        # neither restore semantics nor useful public provenance.
        self.config._name_or_path = ""
        self.config.llm_config._name_or_path = ""
        self.config.llm_config.dtype = get_parameter_dtype(self)
        self.config._sync_llm_surface()
        return super().save_pretrained(*args, **kwargs)

    def prepare_inputs_for_generation(
        self, *args: Any, **kwargs: Any
    ) -> dict[str, Any]:
        return self.llm.prepare_inputs_for_generation(*args, **kwargs)

    def generate_mixed(
        self,
        inputs: "TimeBraidBatchFeature",
        *,
        max_new_tokens: Optional[int] = None,
        max_length: Optional[int] = None,
        mot_target_total_lengths: Optional[Sequence[int]] = None,
        mot_forecast_head_len: Optional[int] = None,
        mot_return_forecast_quantiles: bool = False,
        eos_token_id: Optional[Union[int, Sequence[int]]] = None,
        pad_token_id: Optional[int] = None,
        use_cache: Optional[bool] = None,
        cache_implementation: Optional[str] = None,
    ) -> TimeBraidGenerateOutput:
        """Generate mixed text and time series, or raise.

        `inputs` is the mapping returned by
        `TimeBraidProcessor.apply_chat_template`, which carries `input_ids`,
        `attention_mask`, the complete TS payload, and the target horizons.
        Everything the processor cannot know is a keyword argument here.

        This is the explicit counterpart to `generate`. It states its whole
        contract in its signature, so a misspelled argument is a `TypeError`
        rather than a silent downgrade, and an argument this scheduler cannot
        honour — `temperature`, `logits_processor`, anything sampling-related —
        is simply not a parameter. Unlike `generate`, it never falls back to
        plain text: if the request does not resolve to mixed generation it
        raises and says which signal was missing.

        Supply exactly one of `max_new_tokens` or `max_length`.

        TimeBraid's own arguments keep the `mot_` prefix they carry everywhere
        else — in `config.json`, in the processor's output, and in `generate`'s
        keyword arguments — so one name means one thing at every boundary.
        """
        require_mixed_generation_inputs(inputs)
        input_ids = validate_mot_generation_input_ids(inputs["input_ids"])
        batch_size = int(input_ids.shape[0])

        payload = TimeBraidPayload(
            **{name: inputs.get(name) for name in TIMEBRAID_TS_PAYLOAD_FIELDS}
        )
        controls = NormalizedMoTGenerationControls.normalize(
            batch_size=batch_size,
            # Supplied by the processor, which owns the prompt surface.
            target_horizons=inputs.get("mot_target_horizons"),
            target_history_span_idxs=inputs.get("mot_target_history_span_idxs"),
            # Supplied by the caller: the processor cannot know these.
            target_total_lengths=mot_target_total_lengths,
            forecast_head_len=mot_forecast_head_len,
            return_forecast_quantiles=mot_return_forecast_quantiles,
        )

        # No route resolution here: `require_mixed_generation_inputs` above has
        # already established the one signal that matters, a complete TS
        # payload. Inferring intent is the compatibility shim's problem.
        if controls.target_horizons is None:
            raise RuntimeError(
                "generate_mixed requires `mot_target_horizons` in `inputs`. Pass "
                "`horizon=` to TimeBraidProcessor.apply_chat_template, which "
                "supplies it; use 0 to forbid TS rollout."
            )

        generation_kwargs: dict[str, Any] = {}
        for name, value in (
            ("max_new_tokens", max_new_tokens),
            ("max_length", max_length),
            ("eos_token_id", eos_token_id),
            ("pad_token_id", pad_token_id),
            ("use_cache", use_cache),
            ("cache_implementation", cache_implementation),
        ):
            if value is not None:
                generation_kwargs[name] = value

        return run_timebraid_generate(
            model=self,
            text_model=self.llm.model,
            payload=payload,
            input_ids=input_ids,
            attention_mask=inputs.get("attention_mask"),
            controls=controls,
            kwargs=generation_kwargs,
        )

    def generate(
        self, *args: Any, **kwargs: Any
    ) -> TimeBraidGenerateOutput | torch.Tensor | ModelOutput:
        """Hugging Face-compatible dispatch over `**kwargs`.

        Kept so `model.generate(**processor_output, max_new_tokens=...)` keeps
        working, and so plain-text requests still reach the host language
        model. It resolves a route and delegates; the contract for mixed
        generation lives on `generate_mixed`.
        """
        generation_kwargs = dict(kwargs)
        payload = TimeBraidPayload.pop_from_kwargs(generation_kwargs)
        generation_controls = {
            key: generation_kwargs.pop(key)
            for key in _MOT_GENERATION_CONTROL_FIELDS
            if key in generation_kwargs
        }
        _reject_unknown_timebraid_kwargs(generation_kwargs)

        positional_input = args[0] if args else None
        keyword_input = generation_kwargs.get("input_ids")
        hf_input = generation_kwargs.get("inputs")
        route_input = (
            positional_input
            if positional_input is not None
            else (keyword_input if keyword_input is not None else hf_input)
        )
        if isinstance(route_input, torch.Tensor) and route_input.ndim == 2:
            control_batch_size = int(route_input.shape[0])
        elif isinstance(payload.ts_values, torch.Tensor) and payload.ts_values.ndim > 0:
            control_batch_size = int(payload.ts_values.shape[0])
        else:
            control_batch_size = 1
        controls = normalize_mot_generation_controls(
            generation_controls,
            batch_size=control_batch_size,
        )

        if isinstance(generation_kwargs.get("past_key_values"), MoTDynamicCache):
            raise RuntimeError(
                "MoT native mixed generate accepts a fresh prompt and owns its full "
                "phase scheduler; continue a public forward cache through "
                "model.forward instead."
            )

        decision = resolve_generation_route(
            payload=payload,
            route_input=route_input,
            target_horizons=controls.target_horizons,
            target_total_lengths=controls.target_total_lengths,
            ts_open_token_id=int(self.ts_open_token_id),
            ts_close_token_id=int(self.ts_close_token_id),
        )
        if decision.route is GenerationRoute.PLAIN_TEXT:
            _require_binary_plain_attention_mask(
                generation_kwargs.get("attention_mask")
            )
            return self.llm.generate(*args, **generation_kwargs)
        if controls.target_horizons is None:
            raise RuntimeError(
                "Mixed generation requires explicit `mot_target_horizons`. Pass "
                "`horizon=` to TimeBraidProcessor.apply_chat_template, which "
                "supplies it; use 0 to forbid TS rollout, or a positive integer "
                "to allow one generated TS span."
            )

        if len(args) > 1:
            raise TypeError(
                "TS generation accepts at most one positional input tensor."
            )
        provided_inputs = [
            value
            for value in (positional_input, keyword_input, hf_input)
            if value is not None
        ]
        if len(provided_inputs) != 1:
            raise ValueError(
                "TS generation requires exactly one of positional inputs, "
                "input_ids, or inputs."
            )
        input_ids = validate_mot_generation_input_ids(provided_inputs[0])
        generation_kwargs.pop("input_ids", None)
        generation_kwargs.pop("inputs", None)
        attention_mask = generation_kwargs.pop("attention_mask", None)
        return run_timebraid_generate(
            model=self,
            text_model=self.llm.model,
            payload=payload,
            input_ids=input_ids,
            attention_mask=attention_mask,
            controls=controls,
            kwargs=generation_kwargs,
        )

    def get_input_embeddings(self) -> nn.Module:
        return self.llm.get_input_embeddings()

    def set_input_embeddings(self, value: nn.Module) -> None:
        self.llm.set_input_embeddings(value)

    def get_output_embeddings(self) -> Optional[nn.Module]:
        return self.llm.get_output_embeddings()

    def set_output_embeddings(self, value: nn.Module) -> None:
        self.llm.set_output_embeddings(value)

    def tie_weights(self) -> None:
        super().tie_weights()
        self._sync_hf_loading_contracts()

    def can_generate(self) -> bool:
        llm = self._modules.get("llm") if hasattr(self, "_modules") else None
        return True if llm is None else llm.can_generate()


def register_timebraid_auto_classes() -> None:
    """Register TimeBraid with HF auto classes in the active Python process."""
    from transformers import AutoProcessor

    from ..processing_timebraid import TimeBraidProcessor

    AutoConfig.register(TimeBraidConfig.model_type, TimeBraidConfig, exist_ok=True)
    AutoProcessor.register(TimeBraidConfig, TimeBraidProcessor, exist_ok=True)
    # Deliberately not registered with plain `AutoModel`: that class means
    # "base model without a task head", and TimeBraid is a causal LM. The
    # published auto_map must not advertise it either — see AUTO_MAP in
    # scripts/build-huggingface-timebraid-model.py.
    AutoModelForCausalLM.register(TimeBraidConfig, TimeBraid, exist_ok=True)


__all__ = [
    "TimeBraid",
    "TimeBraidConfig",
    "TimeBraidLossBreakdown",
    "TimeBraidOutput",
    "register_timebraid_auto_classes",
]