model.py

"""LatentQwenASR model for lr_whisper.

The ``LatentQwenASR`` class wraps a frozen Qwen3-ASR base model and adds a
small controller network that emits per-step delta vectors (latent tokens).
"""

import math
import re
import string
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from tqdm import tqdm
from jiwer import wer

from data import build_audio_token_seq, _coerce_feat_attention_mask

def _detach_kv(kv_tuple: Any) -> Any:
    """Detach past_key_values to prevent excessive BPTT exploding gradients."""
    if kv_tuple is None:
        return None
    # If it's a raw tuple (older transformers standard)
    if isinstance(kv_tuple, tuple):
        return tuple(tuple(t.detach() for t in layer) for layer in kv_tuple)

    # Transformers >= 4.38 uses Cache objects like DynamicCache
    if hasattr(kv_tuple, "to_legacy_cache"):
        legacy_tuple = kv_tuple.to_legacy_cache()
        detached_tuple = tuple(tuple(t.detach() for t in layer) for layer in legacy_tuple)
        return type(kv_tuple).from_legacy_cache(detached_tuple)

    # Fallback to cloning attributes if the specific class doesn't match
    import copy
    new_cache = copy.copy(kv_tuple)
    if hasattr(new_cache, "key_cache"):
        new_cache.key_cache = [k.detach() for k in getattr(kv_tuple, "key_cache", [])]
        new_cache.value_cache = [v.detach() for v in getattr(kv_tuple, "value_cache", [])]
    return new_cache


class LatentQwenASR(nn.Module):
    """Wrapper around Qwen3-ASR that supports latent prompt injection.

    The class keeps the underlying ASR model frozen and introduces a small
    controller network that emits delta vectors for each latent token.  These
    deltas are normalised and scaled, then added to the token embeddings at
    the NT-token positions.  During training, the deltas are inserted into the
    decoder embeddings at the appropriate positions; during generation, they
    are applied just before decoding.

    Design note - float32 parameters
    -----------------------------------
    ``log_scale`` is kept in **float32** even when the rest of the model is in bfloat16
    to preserve gradient precision for the per-step scales.
    """

    @staticmethod
    def _resolve_text_model(thinker: nn.Module) -> nn.Module:
        """Resolve the inner text decoder module across wrapped backends (e.g. PEFT)."""
        queue: List[object] = [thinker]
        seen: set = set()
        while queue:
            node = queue.pop(0)
            if node is None:
                continue
            node_id = id(node)
            if node_id in seen:
                continue
            seen.add(node_id)
            emb = getattr(node, "embed_tokens", None)
            if isinstance(emb, nn.Module):
                return node
            for attr in ("model", "base_model", "module"):
                child = getattr(node, attr, None)
                if child is not None and child is not node:
                    queue.append(child)
        raise AttributeError(
            "Unable to resolve text decoder with `embed_tokens` from thinker module."
        )

    @staticmethod
    def _resolve_embed_tokens(text_model: nn.Module) -> nn.Module:
        emb = getattr(text_model, "embed_tokens", None)
        if isinstance(emb, nn.Module):
            return emb
        if hasattr(text_model, "get_input_embeddings"):
            emb = text_model.get_input_embeddings()
            if isinstance(emb, nn.Module):
                return emb
        raise AttributeError("Unable to resolve input embedding module.")

    def __init__(
        self,
        asr_model: nn.Module,
        processor: Any,
        n_latent: int,
        nt_token_id: int,
        lang_token_id: int,
        transcribe_token_id: int,
        freeze_base: bool = True,
        use_latent: bool = True,
        use_soft_prompt: bool = False,
        soft_prompt_init_mode: str = "text",
        soft_prompt_init_text: str = "",
        user_prompt_text: str = "Transcribe the audio into text.",

        delta_tanh_c: float = 5.0,
        scale_max: float = 3.0,
        scale_init: float = 0.2,
        thought_mode: str = "prefix",
        thought_group_size: int = 1,
        halt_threshold: float = 0.0,
        latent_drop_prob: float = 0.0,
        latent_input_noise_std: float = 0.0,
        latent_use_bounded_delta: bool = True,
        latent_use_injection_gate: bool = True,
        latent_use_embedding_anchor: bool = True,
        freeze_audio_stack: bool = True,
    ) -> None:
        super().__init__()
        self.base_model = asr_model
        self.processor = processor
        self.config = getattr(asr_model, "config", None)
        self.n_latent = int(n_latent)
        self.use_latent = bool(use_latent and self.n_latent > 0)
        self.nt_token_id = int(nt_token_id)
        self.use_soft_prompt = bool(
            use_soft_prompt
            and (not self.use_latent)
            and self.n_latent > 0
            and self.nt_token_id >= 0
        )
        self.soft_prompt_init_mode = (soft_prompt_init_mode or "text").strip().lower()
        self.soft_prompt_init_text = soft_prompt_init_text or ""
        self.lang_token_id = int(lang_token_id)
        self.transcribe_token_id = int(transcribe_token_id)

        # Hyperparams stored from constructor args (originally read from env globals).

        self.delta_tanh_c = float(delta_tanh_c)
        self.halt_threshold = float(halt_threshold)
        self.latent_drop_prob = float(latent_drop_prob)
        self.latent_input_noise_std = float(latent_input_noise_std)
        self.latent_use_bounded_delta = bool(latent_use_bounded_delta)
        self.latent_use_injection_gate = bool(latent_use_injection_gate)
        self.latent_use_embedding_anchor = bool(latent_use_embedding_anchor)
        self.freeze_audio_stack = bool(freeze_audio_stack)

        # Qwen3-ASR audio token ID
        _audio_tok = None
        if self.config and hasattr(self.config, "audio_token_id") and self.config.audio_token_id is not None:
            _audio_tok = int(self.config.audio_token_id)
        if _audio_tok is None:
            _thinker_cfg = getattr(self.base_model, "config", None) or getattr(
                getattr(self.base_model, "thinker", None), "config", None
            )
            if _thinker_cfg is not None and hasattr(_thinker_cfg, "audio_token_id") and _thinker_cfg.audio_token_id is not None:
                _audio_tok = int(_thinker_cfg.audio_token_id)
        if _audio_tok is None:
            for _name in ("<|AUDIO|>", "<|audio|>", "<|audio_content|>", "<|audio_pad|>"):
                target_tok = self.processor.tokenizer if hasattr(self.processor, "tokenizer") else self.processor
                if hasattr(target_tok, "convert_tokens_to_ids"):
                    _try = target_tok.convert_tokens_to_ids(_name)
                    unk_id = getattr(target_tok, "unk_token_id", None)
                    if _try is not None and _try != unk_id:
                        _audio_tok = int(_try)
                        break
        if _audio_tok is None:
            _audio_tok = getattr(self.processor, "audio_token_id", None)
        if _audio_tok is None:
            _audio_tok = 151646  # last-resort fallback
        self.audio_token_id = _audio_tok
        
        target_tok = self.processor.tokenizer if hasattr(self.processor, "tokenizer") else self.processor
        _decoded = target_tok.convert_ids_to_tokens(self.audio_token_id) if hasattr(target_tok, "convert_ids_to_tokens") else str(self.audio_token_id)
        print(f"[init] audio_token_id={self.audio_token_id} decodes_to={_decoded!r}")

        # Freeze all parameters in the underlying model for latent-adapter training.
        if freeze_base:
            for p in self.base_model.parameters():
                p.requires_grad = False

        if self.freeze_audio_stack:
            audio_tower = getattr(self.base_model, "thinker", self.base_model)
            audio_tower = getattr(audio_tower, "audio_tower", None)
            if audio_tower is not None:
                n_frozen = 0
                for p in audio_tower.parameters():
                    if p.requires_grad:
                        p.requires_grad = False
                        n_frozen += 1
                if n_frozen > 0:
                    print(f"[init] Froze {n_frozen} audio_tower parameters.")
            # Also freeze multi_modal_projector (audio-to-text projection).
            projector = getattr(getattr(self.base_model, "thinker", self.base_model), "multi_modal_projector", None)
            if projector is not None:
                n_frozen_proj = 0
                for p in projector.parameters():
                    if p.requires_grad:
                        p.requires_grad = False
                        n_frozen_proj += 1
                if n_frozen_proj > 0:
                    print(f"[init] Froze {n_frozen_proj} multi_modal_projector parameters.")

        self.thinker = self.base_model.thinker
        self.text_model = self._resolve_text_model(self.thinker)
        self.embed_tokens = self._resolve_embed_tokens(self.text_model)

        if self.config and hasattr(self.config, "text_config"):
            d = self.config.text_config.hidden_size
        elif self.config and hasattr(self.config, "hidden_size"):
            d = self.config.hidden_size
        else:
            if hasattr(self.embed_tokens, "embedding_dim"):
                d = int(self.embed_tokens.embedding_dim)
            else:
                d = int(self.embed_tokens.weight.size(-1))

        print(f"[LatentQwenASR] Using hidden_size={d}")

        target_dtype = self.thinker.dtype if hasattr(self.thinker, "dtype") else torch.float32

        self.init_proj = nn.Linear(d, d, bias=True).to(dtype=target_dtype)
        self.delta_proj = nn.Linear(d, d, bias=True).to(dtype=target_dtype)

        self.step_embed = nn.Parameter(torch.zeros(self.n_latent, d, dtype=target_dtype))
        nn.init.normal_(self.step_embed, mean=0.0, std=0.02)
        self.step_proj = nn.Linear(d, d, bias=False).to(dtype=target_dtype)
        self.scale_max = float(scale_max)



        scale_init_val = float(scale_init)
        if scale_init_val <= 0:
            scale_init_val = 1.0
        scale_init_unconstrained = math.log(math.expm1(scale_init_val))
        self.log_scale = nn.Parameter(
            torch.full((self.n_latent,), float(scale_init_unconstrained), dtype=torch.float32)
        )

        # Stability: LayerNorm for the recurrent thought loop
        self.thought_ln = nn.LayerNorm(d, elementwise_affine=False).to(dtype=target_dtype)

        # Thought Quality Analyzer (Value Head) predicting expected CE error
        self.value_head = nn.Linear(d, 1).to(dtype=target_dtype)

        # Soft Gated Injection (Scheme 4)
        self.injection_gate = nn.Sequential(
            nn.Linear(d * 2, d),
            nn.Sigmoid()
        ).to(dtype=target_dtype)
        # Initialize gate weights to 0 so sigmoid(0) = 0.5 (neutral start)
        nn.init.zeros_(self.injection_gate[0].weight)
        nn.init.zeros_(self.injection_gate[0].bias)
        self._hidden_size = d

        # Front-token prompt tuning
        self.soft_prompt_embed = nn.Parameter(
            torch.zeros(self.n_latent, d, dtype=target_dtype),
            requires_grad=self.use_soft_prompt,
        )
        if self.use_soft_prompt:
            init_prompt = None
            if self.soft_prompt_init_mode == "text" and self.soft_prompt_init_text.strip():
                init_ids = self.processor.tokenizer.encode(
                    self.soft_prompt_init_text.strip(),
                    add_special_tokens=False,
                )
                if init_ids:
                    with torch.no_grad():
                        emb_device = self.embed_tokens.weight.device
                        init_ids_t = torch.tensor(init_ids, dtype=torch.long, device=emb_device)
                        base = self.embed_tokens(init_ids_t).detach()
                    if base.numel() > 0:
                        if base.size(0) < self.n_latent:
                            reps = (self.n_latent + base.size(0) - 1) // base.size(0)
                            base = base.repeat((reps, 1))
                        init_prompt = base[: self.n_latent]
            if init_prompt is None:
                init_prompt = torch.empty((self.n_latent, d), dtype=target_dtype)
                nn.init.normal_(init_prompt, mean=0.0, std=0.02)
            self.soft_prompt_embed.data.copy_(
                init_prompt.to(device=self.soft_prompt_embed.device, dtype=self.soft_prompt_embed.dtype)
            )
        else:
            self.soft_prompt_embed.requires_grad = False

        # Cast trainable modules to backbone dtype
        target_dtype = self.base_model.dtype if hasattr(self.base_model, "dtype") else torch.float32
        self.init_proj.to(target_dtype)
        self.delta_proj.to(target_dtype)
        self.step_embed.data = self.step_embed.data.to(target_dtype)
        self.step_proj.to(target_dtype)
        self.soft_prompt_embed.data = self.soft_prompt_embed.data.to(target_dtype)
        # log_scale intentionally stays float32 for per-step scale gradient precision.

        # Qwen3-ASR Special Audio Tokens
        self.audio_bos_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|audio_start|>")
        self.audio_eos_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|audio_end|>")
        if self.audio_bos_token_id is None:
            self.audio_bos_token_id = self.processor.tokenizer.bos_token_id
        if self.audio_eos_token_id is None:
            self.audio_eos_token_id = self.processor.tokenizer.eos_token_id

        # Chat-template token pieces
        self.im_start_id = self.processor.tokenizer.convert_tokens_to_ids("<|im_start|>")
        if self.im_start_id is None or self.im_start_id == self.processor.tokenizer.unk_token_id:
            self.im_start_id = self.lang_token_id
        _im_end_id = self.processor.tokenizer.convert_tokens_to_ids("<|im_end|>")
        if _im_end_id is None or _im_end_id == self.processor.tokenizer.unk_token_id:
            _im_end_id = self.processor.tokenizer.eos_token_id
        self.im_end_id = _im_end_id
        # Unified stop-token set for generation AND eval truncation.
        _stop_ids = {int(self.im_end_id)}
        _eot_id = self.processor.tokenizer.convert_tokens_to_ids("<|endoftext|>")
        if _eot_id is not None and _eot_id != self.processor.tokenizer.unk_token_id:
            _stop_ids.add(int(_eot_id))
        if self.processor.tokenizer.eos_token_id is not None:
            _stop_ids.add(int(self.processor.tokenizer.eos_token_id))
        self.stop_ids = sorted(_stop_ids)
        self.user_nl_ids = self.processor.tokenizer.encode("user\n", add_special_tokens=False)
        self.asst_nl_ids = self.processor.tokenizer.encode("assistant\n", add_special_tokens=False)
        self.asr_prefix_ids = self.processor.tokenizer.encode("language English<asr_text>", add_special_tokens=False)
        self.nl_ids = self.processor.tokenizer.encode("\n", add_special_tokens=False)
        self.user_prompt_text = user_prompt_text
        if self.user_prompt_text:
            self.user_prompt_ids = self.processor.tokenizer.encode(
                self.user_prompt_text, add_special_tokens=False
            )
        else:
            self.user_prompt_ids = []
        self._warned_generate_fallback = False
        self.thought_mode = thought_mode.strip().lower()
        self.thought_group_size = max(1, int(thought_group_size))
        print(
            "[latent-ablation] "
            f"bounded_delta={self.latent_use_bounded_delta} "
            f"injection_gate={self.latent_use_injection_gate} "
            f"embedding_anchor={self.latent_use_embedding_anchor}"
        )

        # In non-latent modes, keep latent controller out of optimization.
        if not self.use_latent:
            for p in self.init_proj.parameters():
                p.requires_grad = False
            for p in self.delta_proj.parameters():
                p.requires_grad = False
            for p in self.step_proj.parameters():
                p.requires_grad = False
            for p in self.thought_ln.parameters():
                p.requires_grad = False
            self.step_embed.requires_grad = False

            self.log_scale.requires_grad = False
        if not self.use_soft_prompt:
            self.soft_prompt_embed.requires_grad = False

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _build_chat_template_tensors(
        self,
        device: torch.device,
        language: str = "English",
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Build shared chat-template pieces for train/eval consistency."""
        system_nl_ids = self.processor.tokenizer.encode("system\n", add_special_tokens=False)
        system_turn = torch.tensor(
            [self.im_start_id] + system_nl_ids + [self.im_end_id] + self.nl_ids,
            dtype=torch.long,
            device=device,
        )
        user_prefix = torch.tensor(
            [self.im_start_id] + self.user_nl_ids,
            dtype=torch.long,
            device=device,
        )
        if self.user_prompt_ids:
            user_prompt = torch.tensor(self.user_prompt_ids, dtype=torch.long, device=device)
        else:
            user_prompt = torch.empty((0,), dtype=torch.long, device=device)
        user_suffix = torch.tensor(
            [self.im_end_id] + self.nl_ids,
            dtype=torch.long,
            device=device,
        )
        asr_prefix_ids = self.processor.tokenizer.encode(f"language {language}<asr_text>", add_special_tokens=False)
        assistant_prefix = torch.tensor(
            [self.im_start_id] + self.asst_nl_ids + asr_prefix_ids,
            dtype=torch.long,
            device=device,
        )
        return system_turn, user_prefix, user_prompt, user_suffix, assistant_prefix

    def _build_full_sequence(
        self,
        audio_toks: torch.Tensor,
        text_all: torch.Tensor,
        system_turn: torch.Tensor,
        user_prefix: torch.Tensor,
        user_prompt: torch.Tensor,
        user_suffix: torch.Tensor,
        label_pad: int = -100,
        n_extra_mask: int = 0,
        mask_all_nt: bool = False,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Build full input_ids and aligned labels for one sample.

        The sequence layout is::

            [system_turn, user_prefix, audio_toks, user_prompt,
             user_suffix, text_all]

        Labels clone input_ids, then mask everything up to (and including)
        the assistant control tokens, plus an optional ``n_extra_mask``
        additional positions (e.g. NT tokens in latent prefix mode).

        When ``mask_all_nt=True``, all NT-token positions anywhere in the
        sequence are additionally masked.  Use this for interleaved mode where
        NT tokens are scattered throughout the assistant turn.

        Args:
            audio_toks: 1-D token tensor for the audio segment.
            text_all: 1-D token tensor for the target text (assistant turn).
            system_turn: Shared system-turn token tensor.
            user_prefix: Shared user-prefix token tensor.
            user_prompt: Shared user-prompt token tensor.
            user_suffix: Shared user-suffix token tensor.
            label_pad: Padding value for ignored label positions.
            n_extra_mask: Number of additional positions after the assistant
                prefix to mask (set to ``self.n_latent`` for latent prefix
                mode, 0 for native/baseline/interleaved).
            mask_all_nt: If True, additionally mask every NT-token position
                in the full sequence (for interleaved thought mode).

        Returns:
            ``(full_ids, full_labels)`` - 1-D LongTensors.
        """
        full_ids = torch.cat(
            [system_turn, user_prefix, audio_toks, user_prompt, user_suffix, text_all]
        )
        n_prefix = (
            system_turn.size(0)
            + user_prefix.size(0)
            + audio_toks.size(0)
            + user_prompt.size(0)
            + user_suffix.size(0)
        )
        # Mask: <|im_start|> + "assistant\n" + n_extra_mask (e.g. NT tokens)
        n_asst_control = 1 + len(self.asst_nl_ids) + n_extra_mask
        n_total_mask = n_prefix + n_asst_control
        full_labels = full_ids.clone()
        full_labels[:n_total_mask] = label_pad
        if mask_all_nt and self.nt_token_id >= 0:
            full_labels[full_ids == self.nt_token_id] = label_pad
        return full_ids, full_labels

    def _inject_latent_deltas(
        self,
        inputs_embeds: torch.Tensor,
        input_ids: torch.LongTensor,
        deltas: torch.Tensor,
    ) -> torch.Tensor:
        """Inject latent delta embeddings at NT-token positions.

        Replaces the NT-token positions in *inputs_embeds* with::

            base_nt_embedding + delta

        Deltas are always injected at full strength.  The Value Head controls
        hard routing (N=0 skip / early halt) instead of soft scaling.
        """
        B = inputs_embeds.size(0)
        K = deltas.size(1)
        device = inputs_embeds.device
        nt_mask = (input_ids == self.nt_token_id)
        nt_counts = nt_mask.sum(dim=1)

        max_nt = int(nt_counts.max().item())
        if K == 1 and max_nt > 1:
            deltas = deltas.expand(B, max_nt, -1)
            K = max_nt

        assert (nt_counts == K).all(), (
            f"NT token count mismatch! {nt_counts} vs expected {K}"
        )
        nt_token_ids = torch.full(
            (B, K), self.nt_token_id, dtype=torch.long, device=device
        )
        base_nt = self.embed_tokens(nt_token_ids).to(dtype=inputs_embeds.dtype)
        delta_scaled = deltas.float().to(dtype=inputs_embeds.dtype)

        if self.latent_use_injection_gate:
            gate_input = torch.cat([base_nt, delta_scaled], dim=-1)
            gating_weight = self.injection_gate(gate_input)
        else:
            gating_weight = torch.ones_like(delta_scaled)
        gated_delta = delta_scaled * gating_weight
        if self.latent_use_embedding_anchor:
            nt_embeds = base_nt + gated_delta
        else:
            nt_embeds = gated_delta
        batch_idx, seq_idx = torch.nonzero(nt_mask, as_tuple=True)
        result = inputs_embeds.clone()
        result[batch_idx, seq_idx, :] = nt_embeds.to(dtype=inputs_embeds.dtype).view(-1, self._hidden_size)
        return result

    def _bounded_delta_from_raw(
        self,
        delta_raw: torch.Tensor,
        step_idx: int,
        scales: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Convert a raw projected state into the configured latent delta."""
        if not self.latent_use_bounded_delta:
            return delta_raw
        if scales is None:
            scales = self.step_scales()
        delta_raw_f32 = delta_raw.float()
        norm_val = delta_raw_f32.norm(dim=-1, keepdim=True)
        delta_dir = (delta_raw_f32 / (norm_val + 1e-6)).to(dtype=delta_raw.dtype)
        scale = scales[min(int(step_idx), scales.size(0) - 1)]
        return delta_dir * scale

    def _compute_prefix_state(
        self,
        inputs_embeds: torch.Tensor,
        input_ids: torch.LongTensor,
        strict: bool = True,
        language: str = "English",
    ) -> Tuple[torch.Tensor, object, torch.LongTensor]:
        """Compute prefix state and KV cache up to the first NT token."""
        B, _, _ = inputs_embeds.shape
        device = inputs_embeds.device
        text_model = self.text_model

        nt_mask = (input_ids == self.nt_token_id)
        nt_counts = nt_mask.sum(dim=1)
        if strict:
            assert (nt_counts == self.n_latent).all(), (
                f"NT token count mismatch! {nt_counts} vs {self.n_latent}"
            )

        first_nt = nt_mask.float().argmax(dim=1)
        if strict:
            asr_prefix_ids = self.processor.tokenizer.encode(
                f"language {language}<asr_text>", add_special_tokens=False
            )
            expected = torch.tensor(
                [self.im_start_id] + self.asst_nl_ids + asr_prefix_ids, dtype=torch.long, device=device
            )
            exp_len = expected.numel()
            for i in range(B):
                idx = int(first_nt[i].item())
                start = idx - exp_len
                assert start >= 0, (
                    f"Prompt too short to contain assistant prefix (idx={idx}, exp_len={exp_len})."
                )
                actual = input_ids[i, start:idx]
                assert torch.equal(actual, expected), (
                    "Assistant prefix mismatch before NT. "
                    f"expected={expected.tolist()} actual={actual.tolist()}"
                )

        prefix_lens = first_nt
        if strict:
            assert (prefix_lens > 0).all(), f"Invalid prefix length(s): {prefix_lens}"

        max_prefix = int(prefix_lens.max().item())
        prefix_embeds = inputs_embeds[:, :max_prefix, :]
        prefix_attn = (
            torch.arange(max_prefix, device=device).unsqueeze(0) < prefix_lens.unsqueeze(1)
        ).long()
        if prefix_attn.numel() > 0:
            prefix_embeds = prefix_embeds * prefix_attn.unsqueeze(-1)

        with torch.no_grad():
            out = text_model(
                inputs_embeds=prefix_embeds,
                attention_mask=prefix_attn,
                use_cache=True,
                return_dict=True,
            )
        last_positions = torch.clamp(prefix_lens - 1, min=0)
        batch_idx = torch.arange(B, device=device)
        last = out.last_hidden_state[batch_idx, last_positions].detach()
        return last, out.past_key_values, prefix_attn

    def _embed_scale(self) -> float:
        # Qwen3-ASR does NOT scale embeddings by sqrt(d_model).
        return 1.0


    def step_scales(self) -> torch.Tensor:
        s = F.softplus(self.log_scale)
        s = torch.clamp(s, max=self.scale_max)
        return s[:self.n_latent]

    def forced_decoder_ids_latent(self) -> List[Tuple[int, int]]:
        ids: List[Tuple[int, int]] = [(0, self.im_start_id)]
        cur = 1
        for tid in self.asst_nl_ids:
            ids.append((cur, tid))
            cur += 1
        for i in range(self.n_latent):
            ids.append((cur + i, self.nt_token_id))
        return ids

    def forced_decoder_ids_baseline(self) -> List[Tuple[int, int]]:
        ids: List[Tuple[int, int]] = [(0, self.im_start_id)]
        cur = 1
        for tid in self.asst_nl_ids:
            ids.append((cur, tid))
            cur += 1
        return ids

    def _shift_right(self, labels: torch.LongTensor) -> torch.LongTensor:
        pad_id = self.processor.tokenizer.eos_token_id
        decoder_input_ids = labels.new_full(labels.shape, pad_id)
        decoder_input_ids[:, 0] = (
            self.processor.tokenizer.bos_token_id
            if self.processor.tokenizer.bos_token_id is not None
            else self.lang_token_id
        )
        shifted = labels[:, :-1].clone()
        shifted = shifted.masked_fill(shifted == -100, pad_id)
        decoder_input_ids[:, 1:] = shifted
        return decoder_input_ids

    def _encode_audio(
        self,
        input_features: torch.FloatTensor,
        feature_attention_mask: torch.LongTensor,
    ) -> Tuple[torch.Tensor, List[int]]:
        """Manually encode audio features to get their exact lengths per sample."""
        if feature_attention_mask is not None:
            feature_lens = torch.sum(feature_attention_mask, dim=1)
        else:
            feature_lens = torch.full(
                (input_features.size(0),), input_features.size(2),
                device=input_features.device, dtype=torch.long,
            )

        audio_features = []
        lengths = []
        audio_tower = self.thinker.audio_tower
        projector = getattr(self.thinker, "multi_modal_projector", None)

        for i in range(len(input_features)):
            feat = input_features[i]
            l = feature_lens[i]
            inp = feat[:, :l]
            length_tensor = l.unsqueeze(0)
            out = audio_tower(inp, feature_lens=length_tensor).last_hidden_state
            if projector is not None:
                out = projector(out)
            out = out.squeeze(0)
            audio_features.append(out)
            lengths.append(out.size(0))

        return torch.cat(audio_features, dim=0), lengths

    def _get_audio_lengths(
        self,
        input_features: torch.FloatTensor,
        feature_attention_mask: torch.LongTensor,
    ) -> List[int]:
        """Compute audio output lengths using the exact same formula as the thinker."""
        if feature_attention_mask is not None:
            feature_lens = torch.sum(feature_attention_mask, dim=1)
        else:
            feature_lens = torch.full(
                (input_features.size(0),), input_features.size(2),
                device=input_features.device, dtype=torch.long,
            )

        input_lengths_leave = feature_lens % 100
        feat_lengths = (input_lengths_leave - 1) // 2 + 1
        output_lengths = (
            ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1
            + (feature_lens // 100) * 13
        )
        return output_lengths.tolist()

    def _get_native_audio_embeds(
        self,
        input_ids: torch.LongTensor,
        input_features: torch.FloatTensor,
        feature_attention_mask: torch.LongTensor,
        attention_mask: torch.LongTensor,
    ) -> torch.Tensor:
        """Capture the thinker's native audio-fused embeddings via a forward hook."""
        captured: Dict[str, torch.Tensor] = {}

        def _hook(module: nn.Module, args: Any, kwargs: Any) -> None:
            embeds = kwargs.get("inputs_embeds", None)
            if embeds is None and len(args) > 0:
                embeds = args[0]
            if embeds is not None:
                captured["embeds"] = embeds
            return None

        hook = self.text_model.register_forward_pre_hook(_hook, with_kwargs=True)
        try:
            with torch.no_grad():
                self.thinker(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    input_features=input_features,
                    feature_attention_mask=feature_attention_mask,
                    labels=None,
                    return_dict=True,
                )
        finally:
            hook.remove()

        if "embeds" not in captured:
            raise RuntimeError(
                "_get_native_audio_embeds: forward hook did not capture embeddings. "
                "The thinker's internal forward may use a different code path."
            )
        return captured["embeds"].detach().clone()

    def _compute_continuous_thoughts(
        self,
        initial_state: torch.Tensor,
        prefix_past_key_values: Optional[object] = None,
        prefix_attention_mask: Optional[torch.LongTensor] = None,
        halt_threshold: Optional[float] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, torch.Tensor]]:
        """Compute recurrent latent states + per-step deltas using DEQ fixed-point iteration."""
        B = initial_state.size(0)
        D = initial_state.size(-1)
        device = initial_state.device
        text_model = self.text_model

        initial_norm = initial_state / (initial_state.norm(dim=-1, keepdim=True) + 1e-6)
        h = self.init_proj(initial_norm)

        # Remove random corruption as teacher forcing artificially lowers CE
        _corrupted = False
        
        scales = self.step_scales()
        thought_embeds: List[torch.Tensor] = []
        state_embeds: List[torch.Tensor] = []
        raw_norms: List[torch.Tensor] = []
        scaled_norms: List[torch.Tensor] = []
        v_preds: List[torch.Tensor] = []
        past_kv = prefix_past_key_values
        running_attn = prefix_attention_mask

        # ---- Dynamic Causal Thinking Loop ----
        # Instead of fixed-point iteration (which doesn't converge on frozen LLMs),
        # we generate a sequence of thought tokens.
        max_iter = self.n_latent
        if halt_threshold is None:
            halt_threshold = getattr(self, "halt_threshold", 0.0)

        current_embed = h

        # ---- N=0 Skip Check (Tanh) ----
        # v_pred ∈ [-1,1]: expected delta CE (positive means LR helps).
        # If v_pred < threshold → model is confident LR is NOT needed → skip.
        if not self.training:
            v_init = torch.tanh(self.value_head(h.unsqueeze(1)).squeeze(-1).squeeze(-1))  # (B,)
            if (v_init < halt_threshold).all():
                # Return empty deltas → generate() will remove all NT tokens
                empty_deltas = torch.zeros(B, 0, D, device=device, dtype=h.dtype)
                empty_states = torch.zeros(B, 0, D, device=device, dtype=h.dtype)
                stats = {
                    "predicted_value": v_init.detach(),
                    "deq_iters": torch.tensor(0.0, device=device),
                    "raw_norm_mean": torch.zeros(0, device=device),
                    "raw_norm_std": torch.zeros(0, device=device),
                    "scaled_norm_mean": torch.zeros(0, device=device),
                    "scaled_norm_std": torch.zeros(0, device=device),
                    "cos_mean": torch.zeros(0, device=device),
                    "cos_std": torch.zeros(0, device=device),
                    "scales": scales.detach(),
                    "diff_norm": torch.tensor(0.0, device=device),
                    "step_cos": torch.tensor(0.0, device=device),
                    "v_preds": v_init.unsqueeze(1).detach(),
                    "gate": v_init.unsqueeze(1).detach(),
                    "skipped": True,
                }
                return empty_deltas, empty_states, stats

        with torch.no_grad():
            for k in range(max_iter):
                # The input for step k is step_embed[k] projected, added to the normalized state
                step_idx = min(k, self.step_embed.size(0) - 1)
                step = self.step_proj(self.step_embed[step_idx]).view(1, 1, -1).expand(B, 1, -1)
                
                h_norm = self.thought_ln(current_embed)
                h_input = h_norm.unsqueeze(1) + step

                step_attn = None
                if running_attn is not None:
                    # Causal: attention mask GROWS by 1 at each step
                    one = torch.ones((B, 1), dtype=running_attn.dtype, device=device)
                    step_attn = torch.cat([running_attn, one], dim=1)

                out = text_model(
                    inputs_embeds=h_input,
                    past_key_values=past_kv,
                    attention_mask=step_attn,
                    use_cache=True,
                    return_dict=True,
                )

                # Next state
                h_next = out.last_hidden_state.squeeze(1)
                
                # Compute delta to inject into this token's output representation
                delta_raw = self.delta_proj(h_next)
                delta = self._bounded_delta_from_raw(delta_raw, step_idx, scales)

                # Early Halt Check (Value Head, tanh) — BEFORE committing delta
                v_raw = self.value_head(h_next.unsqueeze(1)).squeeze(-1).squeeze(-1)  # (B,)
                v_pred = torch.tanh(v_raw)

                # Halt when v_pred drops below threshold (model says this step is harmful)
                if not self.training:
                    if (v_pred < halt_threshold).all():
                        break

                # Only commit the delta if we did NOT halt
                thought_embeds.append(delta)
                state_embeds.append(h_next)
                raw_norms.append(delta_raw.norm(dim=-1))
                scaled_norms.append(delta.norm(dim=-1))
                v_preds.append(v_pred)

                # Update context for the *next* iteration
                past_kv = out.past_key_values
                running_attn = step_attn
                current_embed = h_next

            # Grab the past_kv produced exactly prior to the final step
            with torch.no_grad():
                # We need to re-run up to final_step_idx-1 to get the exact past_kv?
                # Actually, no. We can just run from initial_state and re-accumulate.
                # BUT since max_iter <= n_latent is small (e.g. 4), we can just re-run 
                # the *entire* sequence with gradients enabled!
                pass
                
        # To make things simple and correct for training with such a small sequence (N=4),
        # we can just re-run the FULL causal generation with gradients enabled and overwrite.
        # This gives exact gradients through all thinking steps (BPTT).
        if self.training:
            thought_embeds_grad = []
            state_embeds_grad = []
            current_embed_grad = h  # Keep grad flow to init_proj
            running_attn_grad = prefix_attention_mask
            past_kv_grad = prefix_past_key_values
            
            # Re-run for exactly the number of steps we decided to take in the no_grad pass
            for k in range(len(thought_embeds)):
                step_idx = min(k, self.step_embed.size(0) - 1)
                step = self.step_proj(self.step_embed[step_idx]).view(1, 1, -1).expand(B, 1, -1)
                
                h_norm = self.thought_ln(current_embed_grad)
                h_input = h_norm.unsqueeze(1) + step

                step_attn = None
                if running_attn_grad is not None:
                    one = torch.ones((B, 1), dtype=running_attn_grad.dtype, device=device)
                    step_attn = torch.cat([running_attn_grad, one], dim=1)

                out = text_model(
                    inputs_embeds=h_input,
                    past_key_values=past_kv_grad,  # Evolve KV cache like inference!
                    attention_mask=step_attn,
                    use_cache=True,
                    return_dict=True,
                )

                h_next_g = out.last_hidden_state.squeeze(1)
                
                delta_raw_g = self.delta_proj(h_next_g)
                delta_g = self._bounded_delta_from_raw(delta_raw_g, step_idx, scales)

                thought_embeds_grad.append(delta_g)
                state_embeds_grad.append(h_next_g)

                # Detach state so CE-loss gradients don't flow back to previous steps.
                # Each step is independently trained on its own delta (1-step truncated BPTT).
                past_kv_grad = out.past_key_values
                running_attn_grad = step_attn
                current_embed_grad = h_next_g.detach()
                
            # Override the no_grad lists with the gradient-tracked ones
            thought_embeds = thought_embeds_grad
            state_embeds = state_embeds_grad
            
        # Stack sequences — guard against empty list when all steps halted
        if not thought_embeds:
            # All steps rejected by Value Head → return empty (K=0)
            deltas = torch.zeros((B, 0, D), device=device, dtype=h.dtype)
            states = torch.zeros((B, 0, D), device=device, dtype=h.dtype)
            # Use initial state for value prediction
            predicted_value = torch.tanh(self.value_head(h.unsqueeze(1)))
            stats = {
                "predicted_value": predicted_value.squeeze(-1).squeeze(-1),
                "deq_iters": torch.tensor(0.0, device=device),
                "raw_norm_mean": torch.zeros(self.n_latent, device=device),
                "raw_norm_std": torch.zeros(self.n_latent, device=device),
                "scaled_norm_mean": torch.zeros(self.n_latent, device=device),
                "scaled_norm_std": torch.zeros(self.n_latent, device=device),
                "cos_mean": torch.zeros(self.n_latent, device=device),
                "cos_std": torch.zeros(self.n_latent, device=device),
                "scales": scales.detach(),
                "diff_norm": torch.tensor(0.0, device=device),
                "step_cos": torch.tensor(0.0, device=device),
                "v_preds": None,
                "gate": torch.zeros((B, 1), device=device),
                "corrupted": _corrupted,
                "skipped": True,
            }
            return deltas, states, stats

        deltas = torch.stack(thought_embeds, dim=1) # (B, K, D)
        states = torch.stack(state_embeds, dim=1)   # (B, K, D)
        
        # predicted_value: tanh [-1, 1] estimating LR's impact on CE at
        # every latent state. Training supervises each step because inference
        # may halt from any of these value predictions.
        predicted_value = torch.tanh(self.value_head(states)).squeeze(-1)  # (B, K)

        # Gate: average tanh v_pred across all steps, scaled to [0, 1] for gated delta injection
        v_preds_stacked = torch.stack(v_preds, dim=1) if v_preds else None  # (B, K)
        if v_preds_stacked is not None:
            gate = (v_preds_stacked.mean(dim=1, keepdim=True) + 1.0) / 2.0  # (B, 1) — mean confidence mapped to [0, 1]
        else:
            gate = (predicted_value.squeeze(-1) + 1.0) / 2.0  # fallback, scaled to [0, 1]

        deq_iters = len(thought_embeds)

        raw_tensor = torch.stack(raw_norms, dim=1)
        scaled_tensor = torch.stack(scaled_norms, dim=1)
        scales_val = scales.detach()

        with torch.no_grad():
            flat_states = states.view(-1, D)
            flat_states_norm = F.normalize(flat_states, dim=-1)

            word_emb_weight = self.embed_tokens.weight
            vocab_size = word_emb_weight.size(0)
            if vocab_size > 10000:
                idx = torch.randperm(vocab_size, device=device)[:10000]
                word_emb_sub = word_emb_weight[idx]
            else:
                word_emb_sub = word_emb_weight

            word_emb_norm = F.normalize(word_emb_sub, dim=-1)
            sim_matrix = torch.matmul(flat_states_norm, word_emb_norm.t())
            if sim_matrix.numel() > 0:
                max_sim, _ = sim_matrix.max(dim=-1)
                cos_mean_val = max_sim.mean()
                cos_std_val = max_sim.std(unbiased=False)
            else:
                cos_mean_val = flat_states_norm.new_tensor(0.0)
                cos_std_val = flat_states_norm.new_tensor(0.0)

        if deltas.size(1) > 1:
            diffs = deltas[:, 1:] - deltas[:, :-1]
            diff_norm_val = diffs.norm(dim=-1).mean()
            step_cos_val = F.cosine_similarity(deltas[:, 1:], deltas[:, :-1], dim=-1).mean()
        else:
            diff_norm_val = deltas.new_tensor(0.0)
            step_cos_val = deltas.new_tensor(0.0)

        stats = {
            "predicted_value": predicted_value,  # (B, K) for per-step MSE against per-sample target
            "deq_iters": torch.tensor(deq_iters, dtype=torch.float32, device=device),
            "raw_norm_mean": raw_tensor.mean(dim=0).detach() if raw_tensor.ndim > 1 else raw_tensor.detach(),
            "raw_norm_std": raw_tensor.std(dim=0, unbiased=False).detach() if raw_tensor.ndim > 1 else raw_tensor.new_zeros(raw_tensor.size(-1)),
            "scaled_norm_mean": scaled_tensor.mean(dim=0).detach() if scaled_tensor.ndim > 1 else scaled_tensor.detach(),
            "scaled_norm_std": scaled_tensor.std(dim=0, unbiased=False).detach() if scaled_tensor.ndim > 1 else scaled_tensor.new_zeros(scaled_tensor.size(-1)),
            "cos_mean": torch.full((self.n_latent,), cos_mean_val.item(), device=device),
            "cos_std": torch.full((self.n_latent,), cos_std_val.item(), device=device),
            "scales": scales_val,
            "diff_norm": diff_norm_val.detach(),
            "step_cos": step_cos_val.detach(),
            "v_preds": v_preds_stacked.detach() if v_preds_stacked is not None else None,
            "gate": gate.detach(),
            "corrupted": _corrupted,
        }
        return deltas, states, stats

    def _compute_causal_thoughts(
        self,
        inputs_embeds: torch.Tensor,
        input_ids: torch.LongTensor,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Dict[str, torch.Tensor]]:
        """Two-pass causal thought injection for interleaved mode.

        Each NT token at position ``k`` receives a thought delta computed from
        the hidden state at position ``k-1`` (the token immediately before it).
        This makes each thought *causal*: it can attend to all previous context
        including the word generated just before it, which aligns with ASR's
        left-to-right generation structure.

        Pass 1 (``no_grad``): forward the full sequence with **base** NT
        embeddings to capture all hidden states ``h_all``.

        Pass 2 (with grad): for each NT at position k, compute a delta from
        ``h_all[:, k-1, :]`` via ``init_proj → delta_proj → normalize →
        scale``.  Inject the modified NT embedding and return the updated
        ``inputs_embeds``.

        Args:
            inputs_embeds: ``(B, SeqLen, D)`` embedding tensor (base NT
                embeddings already placed at NT positions by
                ``_get_native_audio_embeds``).
            input_ids: ``(B, SeqLen)`` token id tensor to locate NT positions.

        Returns:
            ``(modified_embeds, deltas_padded, states_padded, stats)`` where
            ``deltas_padded`` and ``states_padded`` have shape
            ``(B, max_nt, D)`` (zero-padded to the batch maximum NT count).
        """
        B, SeqLen, D = inputs_embeds.shape
        device = inputs_embeds.device

        nt_mask = (input_ids == self.nt_token_id)  # (B, SeqLen)
        batch_idx, nt_pos = torch.nonzero(nt_mask, as_tuple=True)

        if len(nt_pos) == 0:
            empty_d = inputs_embeds.new_zeros((B, 0, D))
            empty_s = inputs_embeds.new_zeros((B, 0, D))
            dummy = inputs_embeds.new_tensor(0.0)
            stats: Dict[str, torch.Tensor] = {
                "raw_norm_mean": inputs_embeds.new_zeros((0,)),
                "raw_norm_std": inputs_embeds.new_zeros((0,)),
                "scaled_norm_mean": inputs_embeds.new_zeros((0,)),
                "scaled_norm_std": inputs_embeds.new_zeros((0,)),
                "cos_mean": inputs_embeds.new_zeros((0,)),
                "cos_std": inputs_embeds.new_zeros((0,)),
                "scales": inputs_embeds.new_zeros((0,)),
                "diff_norm": dummy,
                "step_cos": dummy,
            }
            return inputs_embeds, empty_d, empty_s, stats

        # --- Pass 1: no_grad forward to capture all hidden states ----------
        with torch.no_grad():
            out1 = self.text_model(
                inputs_embeds=inputs_embeds,
                use_cache=False,
                return_dict=True,
            )
        h_all = out1.last_hidden_state  # (B, SeqLen, D) - detached in no_grad

        # Context for each NT at position k is h_all[:, k-1, :]
        ctx_pos = (nt_pos - 1).clamp(min=0)  # (N_total,)
        contexts = h_all[batch_idx, ctx_pos]  # (N_total, D)

        # --- Pass 2: generate deltas (with grad) ---------------------------
        ctx_norm = F.normalize(contexts.float(), dim=-1).to(dtype=inputs_embeds.dtype)
        h_proj = self.init_proj(ctx_norm)         # (N_total, D)
        delta_raw = self.delta_proj(h_proj)       # (N_total, D)
        if self.latent_use_bounded_delta:
            delta_raw_f32 = delta_raw.float()
            norm_val = delta_raw_f32.norm(dim=-1, keepdim=True)
            delta_dir = (delta_raw_f32 / (norm_val + 1e-6)).to(dtype=delta_raw.dtype)
            scale = self.step_scales().mean()     # scalar - shared across all positions
            delta = delta_dir * scale             # (N_total, D)
        else:
            delta = delta_raw

        # Build NT embeddings and inject
        # base_nt_ids and nt_embeds are moved down to use the scaled delta


        # --- Pack into (B, max_nt, D) for loss functions -------------------
        nt_counts = nt_mask.sum(dim=1)  # (B,)
        max_nt = int(nt_counts.max().item())
        deltas_padded = inputs_embeds.new_zeros((B, max_nt, D))
        states_padded = inputs_embeds.new_zeros((B, max_nt, D))
        for b in range(B):
            b_mask = (batch_idx == b)
            cnt = int(nt_counts[b].item())
            if cnt > 0:
                deltas_padded[b, :cnt] = delta[b_mask]
                states_padded[b, :cnt] = contexts[b_mask]

        # --- Stats for logging (shape matches prefix-mode stats) ------------
        with torch.no_grad():
            delta_for_stats = delta.float()
            raw_norms = delta_raw.norm(dim=-1)       # (N_total,)
            scaled_norms = delta_for_stats.norm(dim=-1)
            flat_ctx_norm = F.normalize(contexts, dim=-1)
            word_emb = self.embed_tokens.weight
            vocab_size = word_emb.size(0)
            if vocab_size > 10000:
                idx = torch.randperm(vocab_size, device=device)[:10000]
                word_emb_sub = word_emb[idx]
            else:
                word_emb_sub = word_emb
            word_emb_norm = F.normalize(word_emb_sub, dim=-1)
            sim_matrix = torch.matmul(flat_ctx_norm.to(word_emb_norm.dtype), word_emb_norm.t())
            max_sim, _ = sim_matrix.max(dim=-1)
            cos_mean_val = max_sim.mean()
            cos_std_val = max_sim.std()
            # Compute diff_norm and step_cos across padded deltas
            if max_nt > 1:
                diff = deltas_padded[:, 1:, :] - deltas_padded[:, :-1, :]
                diff_norm_val = diff.norm(dim=-1).mean()
                step_cos_val = F.cosine_similarity(
                    deltas_padded[:, 1:, :], deltas_padded[:, :-1, :], dim=-1
                ).mean()
            else:
                diff_norm_val = deltas_padded.new_tensor(0.0)
                step_cos_val = deltas_padded.new_tensor(0.0)

        # Inject deltas into embeddings (grad-tracked for training)
        base_nt_ids = torch.full(
            (len(batch_idx),), self.nt_token_id, dtype=torch.long, device=device
        )
        base_nt = self.embed_tokens(base_nt_ids).to(dtype=inputs_embeds.dtype)
        delta_i = delta.to(dtype=inputs_embeds.dtype)
        if self.latent_use_injection_gate:
            gate_input = torch.cat([base_nt, delta_i], dim=-1)
            delta_i = delta_i * self.injection_gate(gate_input)
        if self.latent_use_embedding_anchor:
            nt_embeds = base_nt + delta_i
        else:
            nt_embeds = delta_i

        result_embeds = inputs_embeds.clone()
        result_embeds[batch_idx, nt_pos] = nt_embeds

        n_out = max_nt
        stats = {
            "raw_norm_mean": raw_norms.mean().unsqueeze(0).expand(n_out).detach(),
            "raw_norm_std": raw_norms.std(unbiased=False).unsqueeze(0).expand(n_out).detach(),
            "scaled_norm_mean": scaled_norms.mean().unsqueeze(0).expand(n_out).detach(),
            "scaled_norm_std": scaled_norms.std(unbiased=False).unsqueeze(0).expand(n_out).detach(),
            "cos_mean": torch.full((n_out,), cos_mean_val.item(), device=device),
            "cos_std": torch.full((n_out,), cos_std_val.item(), device=device),
            "scales": scale.unsqueeze(0).expand(n_out).detach(),
            "diff_norm": diff_norm_val.detach(),
            "step_cos": step_cos_val.detach(),
        }
        return result_embeds, deltas_padded, states_padded, stats

    @torch.no_grad()
    def _generate_interleaved(
        self,
        base_embeds: torch.Tensor,
        attention_mask: torch.LongTensor,
        max_new_tokens: int,
        do_sample: bool = False,
        temperature: float = 1.0,
    ) -> torch.LongTensor:
        """Token-by-token generation with interleaved thought injection.

        Before each group of ``self.thought_group_size`` generated word tokens,
        a thought delta is computed from the current sequence's last hidden
        state, and an NT token (with the modified embedding) is prepended.

        NT tokens in the output are included as ``self.nt_token_id`` so that
        callers using ``skip_special_tokens=True`` in ``tokenizer.decode`` will
        strip them automatically (since ``<|latent|>`` is a registered special
        token).

        Args:
            base_embeds: ``(1, SeqLen, D)`` audio-fused embeddings of the
                prompt prefix (no NT tokens appended yet).
            attention_mask: ``(1, SeqLen)`` attention mask for *base_embeds*.
            max_new_tokens: Maximum number of word tokens to generate.
            do_sample: Whether to sample instead of greedy decode.
            temperature: Sampling temperature (ignored if ``do_sample=False``).

        Returns:
            ``(1, N_generated)`` LongTensor of token ids including interleaved
            NT tokens and ending with the first EOS/stop token encountered.
        """
        eos_ids: set = set(self.stop_ids)
        device = base_embeds.device
        safe_temp = float(temperature) if float(temperature) > 0 else 1.0

        cur_embeds = base_embeds.clone()
        cur_attn = attention_mask.clone()
        generated: List[torch.Tensor] = []
        words_generated = 0

        nt_id_t = torch.tensor([[self.nt_token_id]], dtype=torch.long, device=device)
        base_nt_emb = self.embed_tokens(nt_id_t).to(dtype=cur_embeds.dtype)  # (1, 1, D)

        def _append_thought() -> None:
            """Compute a thought from the last hidden state and inject NT."""
            nonlocal cur_embeds, cur_attn
            # Get context = last hidden state of current sequence
            ctx_out = self.text_model(
                inputs_embeds=cur_embeds,
                use_cache=False,
                return_dict=True,
            )
            h_last = ctx_out.last_hidden_state[:, -1:, :]  # (1, 1, D)
            h_norm = F.normalize(h_last.float().squeeze(1), dim=-1).to(dtype=cur_embeds.dtype)
            h_proj = self.init_proj(h_norm)
            delta_raw = self.delta_proj(h_proj)
            if self.latent_use_bounded_delta:
                delta_dir = F.normalize(delta_raw, dim=-1)
                scale = self.step_scales().mean()
                delta_scaled = (delta_dir * scale).float().to(dtype=cur_embeds.dtype)
            else:
                delta_scaled = delta_raw.float().to(dtype=cur_embeds.dtype)
            
            # Use Tanh Value Head for Interleaved mode as well
            if hasattr(self, "value_head"):
                v_raw = self.value_head(h_norm) # (1, 1)
                v_pred = (torch.tanh(v_raw) + 1.0) / 2.0 # (1, 1) -> scaled to [0, 1]
                halt_thresh = float(getattr(self.config, "halt_threshold", 0.0))
                # Skip injecting NT if predicted value < threshold
                if v_pred.squeeze().item() < halt_thresh:
                    return

            delta_i = delta_scaled.unsqueeze(1)
            if self.latent_use_injection_gate:
                gate_input = torch.cat([base_nt_emb.expand_as(delta_i), delta_i], dim=-1)
                delta_i = delta_i * self.injection_gate(gate_input)
            if self.latent_use_embedding_anchor:
                nt_emb_mod = base_nt_emb + delta_i
            else:
                nt_emb_mod = delta_i
            cur_embeds = torch.cat([cur_embeds, nt_emb_mod], dim=1)
            cur_attn = torch.cat([cur_attn, cur_attn.new_ones((1, 1))], dim=1)
            generated.append(nt_id_t.clone())

        for _ in range(max_new_tokens + len(generated)):
            # Inject thought before each new group of words
            if words_generated % self.thought_group_size == 0:
                _append_thought()

            # Generate next word token
            out = self.thinker(
                inputs_embeds=cur_embeds,
                attention_mask=cur_attn,
                input_features=None,
                use_cache=False,
                return_dict=True,
            )
            next_logits = out.logits[:, -1, :]
            if do_sample:
                probs = F.softmax(next_logits / safe_temp, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1)
            else:
                next_token = torch.argmax(next_logits, dim=-1, keepdim=True)

            tok_id = int(next_token[0, 0].item())
            generated.append(next_token)
            if tok_id in eos_ids:
                break

            next_emb = self.embed_tokens(next_token)
            cur_embeds = torch.cat([cur_embeds, next_emb], dim=1)
            cur_attn = torch.cat([cur_attn, cur_attn.new_ones((1, 1))], dim=1)
            words_generated += 1

        if not generated:
            return torch.zeros((1, 0), dtype=torch.long, device=device)
        return torch.cat(generated, dim=1)

    # ------------------------------------------------------------------
    # Forward passes
    # ------------------------------------------------------------------

    def _forward_native(
        self,
        input_features: torch.FloatTensor,
        labels: torch.LongTensor,
        feature_attention_mask: torch.LongTensor,
        language_hint: Optional[str] = None,
    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Native thinker forward for LoRA / baseline FT."""
        B = input_features.size(0)
        device = input_features.device

        pad_token_id = self.processor.tokenizer.pad_token_id
        if pad_token_id is None:
            pad_token_id = self.processor.tokenizer.eos_token_id or 151643

        input_ids_list = []
        labels_list = []
        label_pad = -100

        audio_lengths = self._get_audio_lengths(input_features, feature_attention_mask)
        system_turn, user_prefix, user_prompt, user_suffix, _ = (
            self._build_chat_template_tensors(device, language=language_hint or "English")
        )

        for i in range(B):
            valid_len = (labels[i] != label_pad).sum().item()
            text_all = labels[i, :valid_len]
            l_audio = audio_lengths[i]

            audio_toks = build_audio_token_seq(
                self.audio_bos_token_id, self.audio_token_id, self.audio_eos_token_id,
                l_audio, device,
            )
            full_ids, full_labels = self._build_full_sequence(
                audio_toks, text_all, system_turn, user_prefix, user_prompt, user_suffix,
                label_pad=label_pad, n_extra_mask=0,
            )
            input_ids_list.append(full_ids)
            labels_list.append(full_labels)

        max_len = max(x.size(0) for x in input_ids_list)
        input_ids_batch = torch.full((B, max_len), pad_token_id, dtype=torch.long, device=device)
        labels_batch = torch.full((B, max_len), label_pad, dtype=torch.long, device=device)
        for i in range(B):
            l = input_ids_list[i].size(0)
            input_ids_batch[i, :l] = input_ids_list[i]
            labels_batch[i, :l] = labels_list[i]

        attention_mask = (input_ids_batch != pad_token_id).long()

        out = self.thinker(
            input_ids=input_ids_batch,
            attention_mask=attention_mask,
            input_features=input_features,
            feature_attention_mask=feature_attention_mask,
            labels=None,
            return_dict=True,
        )
        logits = out.logits

        # One-time audio diagnostic
        if not getattr(self, "_native_diag_done", False):
            self._native_diag_done = True
            import sys as _s
            _d = _s.stderr.write
            _d(f"\n[native-diag] thinker type: {type(self.thinker).__name__}\n")
            _d(f"[native-diag] has base_model (PEFT): {hasattr(self.thinker, 'base_model')}\n")
            n_audio_ph = (input_ids_batch[0] == self.audio_token_id).sum().item()
            feat_valid = (
                feature_attention_mask[0].sum().item()
                if feature_attention_mask is not None
                else input_features.size(2)
            )
            _d(f"[native-diag] audio_token_id={self.audio_token_id} decoded={self.processor.tokenizer.convert_ids_to_tokens(self.audio_token_id)!r}\n")
            _d(f"[native-diag] input_features: {list(input_features.shape)}, valid_mel_frames={feat_valid}\n")
            _d(f"[native-diag] input_ids shape: {list(input_ids_batch.shape)}\n")
            _d(f"[native-diag] input_features shape: {list(input_features.shape)}\n")
            _d(f"[native-diag] audio placeholder tokens in input_ids[0]: {n_audio_ph}\n")
            _d(f"[native-diag] audio_bos/eos in ids[0]: bos={int((input_ids_batch[0]==self.audio_bos_token_id).sum())} eos={int((input_ids_batch[0]==self.audio_eos_token_id).sum())}\n")
            with torch.no_grad():
                try:
                    zero_feat = torch.zeros_like(input_features[:1])
                    out_noaud = self.thinker(
                        input_ids=input_ids_batch[:1],
                        attention_mask=attention_mask[:1],
                        input_features=zero_feat,
                        feature_attention_mask=feature_attention_mask[:1] if feature_attention_mask is not None else None,
                        labels=None,
                        return_dict=True,
                    )
                    diff = (logits[0] - out_noaud.logits[0]).abs().mean().item()
                    _d(f"[native-diag] logit diff (real vs zero audio): {diff:.6f}\n")
                    if diff < 0.01:
                        _d(f"[native-diag] ⚠️  AUDIO NOT BEING USED! Logits identical with zero audio.\n")
                    else:
                        _d(f"[native-diag] ✓ Audio IS affecting logits (diff={diff:.4f})\n")
                except Exception as e:
                    _d(f"[native-diag] zero-audio test failed: {e}\n")
            tok_str = self.processor.tokenizer.decode(input_ids_batch[0, :15].tolist())
            _d(f"[native-diag] input_ids[0][:15] decoded: {repr(tok_str)}\n")
            _s.stderr.flush()

        # Pre-shift labels for loss
        labels_for_loss = torch.full_like(labels_batch, label_pad)
        labels_for_loss[:, :-1] = labels_batch[:, 1:]

        D = self._hidden_size
        dummy = logits.new_tensor(0.0)
        stats = {
            "raw_norm_mean": logits.new_zeros((0,)),
            "raw_norm_std": logits.new_zeros((0,)),
            "scaled_norm_mean": logits.new_zeros((0,)),
            "scaled_norm_std": logits.new_zeros((0,)),
            "cos_mean": logits.new_zeros((0,)),
            "cos_std": logits.new_zeros((0,)),
            "scales": logits.new_zeros((0,)),
            "diff_norm": dummy,
            "step_cos": dummy,
        }
        deltas = logits.new_zeros((B, 0, D))
        states = logits.new_zeros((B, 0, D))
        initial_state = logits.new_zeros((B, D))

        return logits, stats, deltas, states, labels_for_loss, initial_state

    def forward(
        self,
        input_features: torch.FloatTensor,
        labels: torch.LongTensor,
        feature_attention_mask: Optional[torch.LongTensor] = None,
        global_step: int = 0,
        language_hint: Optional[str] = None,
    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Forward pass for training."""
        B = input_features.size(0)
        device = input_features.device

        if feature_attention_mask is None:
            feature_attention_mask = torch.ones(
                (B, input_features.size(2)), dtype=torch.long, device=device
            )

        # Fast path: use thinker's native forward for LoRA / baseline FT.
        if not self.use_latent and not self.use_soft_prompt:
            return self._forward_native(input_features, labels, feature_attention_mask, language_hint=language_hint)

        audio_lengths = self._get_audio_lengths(input_features, feature_attention_mask)

        new_input_ids_list = []
        new_labels_list = []

        label_pad = -100
        pad_token_id = self.processor.tokenizer.pad_token_id
        if pad_token_id is None:
            pad_token_id = self.processor.tokenizer.eos_token_id or 151643

        system_turn, user_prefix, user_prompt, user_suffix, _ = (
            self._build_chat_template_tensors(device, language=language_hint or "English")
        )

        interleaved_mode = self.use_latent and self.thought_mode == "interleaved"

        # +++ LATENT DROP +++
        drop_latent_this_batch = False
        if self.training and self.use_latent and not interleaved_mode:
            drop_prob = self.latent_drop_prob
            if drop_prob > 0.0 and torch.rand(1).item() < drop_prob:
                drop_latent_this_batch = True
        
        actual_n_latent = 0 if drop_latent_this_batch else self.n_latent

        for i in range(B):
            valid_len = (labels[i] != label_pad).sum().item()
            text_all = labels[i, :valid_len]
            if drop_latent_this_batch:
                text_all = text_all[text_all != self.nt_token_id]

            l_audio = audio_lengths[i]

            audio_toks = build_audio_token_seq(
                self.audio_bos_token_id, self.audio_token_id, self.audio_eos_token_id,
                l_audio, device,
            )
            if interleaved_mode:
                # NT tokens are scattered throughout text_all; mask all NT positions.
                full_ids, full_labels = self._build_full_sequence(
                    audio_toks, text_all, system_turn, user_prefix, user_prompt, user_suffix,
                    label_pad=label_pad, n_extra_mask=0, mask_all_nt=True,
                )
            else:
                # Prefix mode: all NT tokens appear contiguously after asst_nl.
                full_ids, full_labels = self._build_full_sequence(
                    audio_toks, text_all, system_turn, user_prefix, user_prompt, user_suffix,
                    label_pad=label_pad, n_extra_mask=actual_n_latent,
                )
            new_input_ids_list.append(full_ids)
            new_labels_list.append(full_labels)

        max_len = max(x.size(0) for x in new_input_ids_list)
        new_input_ids = torch.full((B, max_len), pad_token_id, dtype=torch.long, device=device)
        new_labels_raw = torch.full((B, max_len), label_pad, dtype=torch.long, device=device)

        for i in range(B):
            l = new_input_ids_list[i].size(0)
            new_input_ids[i, :l] = new_input_ids_list[i]
            new_labels_raw[i, :l] = new_labels_list[i]

        attention_mask = (new_input_ids != pad_token_id).long()

        inputs_embeds = self._get_native_audio_embeds(
            input_ids=new_input_ids,
            input_features=input_features,
            feature_attention_mask=feature_attention_mask,
            attention_mask=attention_mask,
        )

        # +++ LATENT INPUT NOISE (Value Head Regularization) +++
        # Save pre-LR embeddings (with noise) for fair baseline CE comparison
        pre_lr_embeds = None
        if self.training:
            noise_std = self.latent_input_noise_std
            if noise_std > 0.0:
                noise = torch.randn_like(inputs_embeds) * noise_std
                inputs_embeds = inputs_embeds + noise
            # Save the (possibly noisy) embeddings BEFORE LR injection
            if self.use_latent:
                pre_lr_embeds = inputs_embeds.detach().clone()

        # Latent / soft-prompt injection
        if self.use_latent and interleaved_mode:
            # Interleaved: two-pass causal thought generation.
            inputs_embeds, deltas, states, stats = self._compute_causal_thoughts(
                inputs_embeds, new_input_ids
            )
            initial_state = inputs_embeds.new_zeros((B, self._hidden_size))
        elif self.use_latent:
            if actual_n_latent == 0:
                initial_state = inputs_embeds.new_zeros((B, self._hidden_size))
                deltas = inputs_embeds.new_zeros((B, 0, self._hidden_size))
                states = inputs_embeds.new_zeros((B, 0, self._hidden_size))
                stats = {
                    "raw_norm_mean": inputs_embeds.new_zeros((0,)),
                    "raw_norm_std": inputs_embeds.new_zeros((0,)),
                    "scaled_norm_mean": inputs_embeds.new_zeros((0,)),
                    "scaled_norm_std": inputs_embeds.new_zeros((0,)),
                    "cos_mean": inputs_embeds.new_zeros((0,)),
                    "cos_std": inputs_embeds.new_zeros((0,)),
                    "scales": inputs_embeds.new_zeros((0,)),
                    "diff_norm": inputs_embeds.new_tensor(0.0),
                    "step_cos": inputs_embeds.new_tensor(0.0),
                }
            else:
                # Prefix mode: recurrent thought loop from audio-prefix state.
                initial_state, prefix_past_kv, prefix_attn = self._compute_prefix_state(
                    inputs_embeds, new_input_ids, strict=True
                )
                deltas, states, stats = self._compute_continuous_thoughts(
                    initial_state,
                    prefix_past_key_values=prefix_past_kv,
                    prefix_attention_mask=prefix_attn,
                )
                inputs_embeds = self._inject_latent_deltas(
                    inputs_embeds, new_input_ids, deltas
                )
        elif self.use_soft_prompt:
            nt_mask = (new_input_ids == self.nt_token_id)
            nt_counts = nt_mask.sum(dim=1)
            assert (nt_counts == self.n_latent).all(), (
                f"NT token count mismatch! {nt_counts} vs {self.n_latent}"
            )
            prompt = self.soft_prompt_embed.to(dtype=inputs_embeds.dtype)
            prompt = prompt.unsqueeze(0).expand(B, -1, -1).contiguous()
            batch_idx, seq_idx = torch.nonzero(nt_mask, as_tuple=True)
            prompt_embeds = inputs_embeds.clone()
            prompt_embeds[batch_idx, seq_idx, :] = prompt.view(-1, self._hidden_size)
            inputs_embeds = prompt_embeds
            initial_state = inputs_embeds.new_zeros((B, self._hidden_size))
            deltas = inputs_embeds.new_zeros((B, 0, self._hidden_size))
            states = inputs_embeds.new_zeros((B, 0, self._hidden_size))
            stats = {
                "raw_norm_mean": inputs_embeds.new_zeros((0,)),
                "raw_norm_std": inputs_embeds.new_zeros((0,)),
                "scaled_norm_mean": inputs_embeds.new_zeros((0,)),
                "scaled_norm_std": inputs_embeds.new_zeros((0,)),
                "cos_mean": inputs_embeds.new_zeros((0,)),
                "cos_std": inputs_embeds.new_zeros((0,)),
                "scales": inputs_embeds.new_zeros((0,)),
                "diff_norm": inputs_embeds.new_tensor(0.0),
                "step_cos": inputs_embeds.new_tensor(0.0),
            }
        else:
            initial_state = inputs_embeds.new_zeros((B, self._hidden_size))
            deltas = inputs_embeds.new_zeros((B, 0, self._hidden_size))
            states = inputs_embeds.new_zeros((B, 0, self._hidden_size))
            stats = {
                "raw_norm_mean": inputs_embeds.new_zeros((0,)),
                "raw_norm_std": inputs_embeds.new_zeros((0,)),
                "scaled_norm_mean": inputs_embeds.new_zeros((0,)),
                "scaled_norm_std": inputs_embeds.new_zeros((0,)),
                "cos_mean": inputs_embeds.new_zeros((0,)),
                "cos_std": inputs_embeds.new_zeros((0,)),
                "scales": inputs_embeds.new_zeros((0,)),
                "diff_norm": inputs_embeds.new_tensor(0.0),
                "step_cos": inputs_embeds.new_tensor(0.0),
            }

        out = self.thinker(
            inputs_embeds=inputs_embeds,
            input_ids=new_input_ids,
            attention_mask=attention_mask,
            input_features=None,
            feature_attention_mask=None,
            labels=None,
            use_cache=False,
            return_dict=True,
        )
        logits = out.logits

        new_labels = torch.full_like(new_labels_raw, label_pad)
        new_labels[:, :-1] = new_labels_raw[:, 1:]

        if logits.size(1) != new_labels.size(1):
            if logits.size(1) > new_labels.size(1):
                shift = logits.size(1) - new_labels.size(1)
                logits = logits[:, shift:, :]
            else:
                new_labels = new_labels[:, -logits.size(1):]

        # ---- Baseline CE vs LR CE for Value Head MSE (Delta CE) ----
        # Compute per-sample CE for both baseline (no LR) and LR-enhanced
        # forward passes. The difference (baseline_ce - lr_ce) forms the
        # continuous target for the Value Head via tanh compression.
        if self.use_latent and self.training and deltas.size(1) > 0:
            with torch.no_grad():
                # Use the SAME (possibly noisy) embeddings but WITHOUT LR deltas
                # This ensures a fair comparison: noisy-no-LR vs noisy-with-LR
                if pre_lr_embeds is not None:
                    baseline_embeds = pre_lr_embeds
                else:
                    baseline_embeds = self._get_native_audio_embeds(
                        input_ids=new_input_ids,
                        input_features=input_features,
                        feature_attention_mask=feature_attention_mask,
                        attention_mask=attention_mask,
                    )
                baseline_out = self.thinker(
                    inputs_embeds=baseline_embeds,
                    input_ids=new_input_ids,
                    attention_mask=attention_mask,
                    input_features=None,
                    feature_attention_mask=None,
                    labels=None,
                    use_cache=False,
                    return_dict=True,
                )
                baseline_logits = baseline_out.logits
                if baseline_logits.size(1) != new_labels.size(1):
                    if baseline_logits.size(1) > new_labels.size(1):
                        shift = baseline_logits.size(1) - new_labels.size(1)
                        baseline_logits = baseline_logits[:, shift:, :]
                    else:
                        bl_labels = new_labels[:, -baseline_logits.size(1):]
                else:
                    bl_labels = new_labels

                # Per-SAMPLE CE: for each sample in batch, compute CE on valid tokens.
                # This is used by train.py to form the continuous delta_ce target.
                valid_mask = (bl_labels != -100)  # (B, T)
                valid_per_sample = valid_mask.float().sum(dim=1).clamp(min=1)  # (B,)

                # For value head accuracy: exclude asr_prefix tokens
                # ("language", "English", "<asr_text>") which are trivially correct
                # and dilute the real transcript accuracy signal.
                asr_prefix_set = set(self.asr_prefix_ids)
                asr_prefix_mask = torch.zeros_like(bl_labels, dtype=torch.bool)
                for tid in asr_prefix_set:
                    asr_prefix_mask |= (bl_labels == tid)
                acc_mask = valid_mask & ~asr_prefix_mask  # exclude prefix tokens
                acc_per_sample = acc_mask.float().sum(dim=1).clamp(min=1)

                if not hasattr(self, '_acc_debug_printed'):
                    self._acc_debug_printed = True
                    total_tokens = bl_labels.size(1)
                    ce_valid = valid_per_sample.mean().item()
                    acc_valid = acc_per_sample.mean().item()
                    sample_acc_toks = bl_labels[0][acc_mask[0]][:20].tolist()
                    sample_decoded = self.processor.tokenizer.decode(sample_acc_toks)
                    print(f"\n[ACC-DEBUG] total={total_tokens} ce_valid={ce_valid:.0f} acc_valid={acc_valid:.0f}")
                    print(f"[ACC-DEBUG] first 20 acc label ids: {sample_acc_toks}")
                    print(f"[ACC-DEBUG] decoded: {sample_decoded}")
                if valid_mask.any():
                    # Per-sample baseline CE (uses full valid_mask for loss)
                    loss_fct_none = torch.nn.CrossEntropyLoss(ignore_index=-100, reduction='none')
                    bl_ce_per_token = loss_fct_none(
                        baseline_logits.reshape(-1, baseline_logits.size(-1)),
                        bl_labels.reshape(-1),
                    ).view(B, -1)  # (B, T)
                    baseline_ce_per_sample = bl_ce_per_token.sum(dim=1) / valid_per_sample  # (B,)
                    stats["baseline_ce"] = baseline_ce_per_sample  # (B,)

                    # Per-sample LR CE (for fair Delta CE comparison)
                    lr_logits_for_ce = logits
                    if lr_logits_for_ce.size(1) != bl_labels.size(1):
                        lr_logits_for_ce = lr_logits_for_ce[:, -bl_labels.size(1):, :]
                    lr_ce_per_token = loss_fct_none(
                        lr_logits_for_ce.reshape(-1, lr_logits_for_ce.size(-1)),
                        bl_labels.reshape(-1),
                    ).view(B, -1)
                    lr_ce_per_sample = lr_ce_per_token.sum(dim=1) / valid_per_sample
                    stats["lr_ce"] = lr_ce_per_sample # (B,)

                    # Per-sample accuracy for VALUE HEAD (uses acc_mask, excludes prefix)
                    baseline_preds = baseline_logits.argmax(dim=-1)  # (B, T)
                    lr_preds = logits.argmax(dim=-1)                 # (B, T)
                    if lr_preds.size(1) != bl_labels.size(1):
                        lr_preds = lr_preds[:, -bl_labels.size(1):]
                    baseline_correct = ((baseline_preds == bl_labels) & acc_mask).float()
                    lr_correct = ((lr_preds == bl_labels) & acc_mask).float()
                    baseline_acc = baseline_correct.sum(dim=1) / acc_per_sample  # (B,)
                    lr_acc = lr_correct.sum(dim=1) / acc_per_sample              # (B,)
                    stats["baseline_acc"] = baseline_acc
                    stats["lr_acc"] = lr_acc

                    # Error-conditional metrics for value head target (uses acc_mask)
                    baseline_wrong = ((baseline_preds != bl_labels) & acc_mask)  # (B, T)
                    baseline_right = ((baseline_preds == bl_labels) & acc_mask)  # (B, T)
                    lr_fixes = ((lr_preds == bl_labels) & baseline_wrong).float().sum(dim=1)   # (B,)
                    lr_breaks = ((lr_preds != bl_labels) & baseline_right).float().sum(dim=1)  # (B,)
                    baseline_errors = baseline_wrong.float().sum(dim=1).clamp(min=1)           # (B,)
                    stats["lr_fixes"] = lr_fixes
                    stats["lr_breaks"] = lr_breaks
                    stats["baseline_errors"] = baseline_errors
                else:
                    B_size = logits.size(0)
                    stats["baseline_ce"] = logits.new_zeros(B_size)
                    stats["baseline_acc"] = logits.new_ones(B_size)
                    stats["lr_acc"] = logits.new_ones(B_size)
                    stats["lr_fixes"] = logits.new_zeros(B_size)
                    stats["lr_breaks"] = logits.new_zeros(B_size)
                    stats["baseline_errors"] = logits.new_ones(B_size)

        return logits, stats, deltas, states, new_labels, initial_state

    @torch.no_grad()
    def _generate_with_forward_fallback(
        self,
        start_ids: torch.LongTensor,
        inputs_embeds: torch.FloatTensor,
        attention_mask: torch.LongTensor,
        max_new_tokens: int,
        eos_token_id: Optional[Any],
        do_sample: bool,
        temperature: float,
    ) -> torch.LongTensor:
        """Fallback decoder when ``.generate(inputs_embeds=...)`` is unsupported."""
        if eos_token_id is None:
            eos_ids: List[int] = []
        elif isinstance(eos_token_id, (list, tuple, set)):
            eos_ids = [int(x) for x in eos_token_id]
        else:
            eos_ids = [int(eos_token_id)]

        generated_tokens: List[torch.Tensor] = []
        safe_temp = float(temperature) if float(temperature) > 0 else 1.0

        cur_ids = start_ids.clone()
        cur_embeds = inputs_embeds.clone()
        cur_attn = attention_mask.clone()

        for _ in range(int(max_new_tokens)):
            out = self.thinker(
                inputs_embeds=cur_embeds,
                attention_mask=cur_attn,
                input_features=None,
                use_cache=False,
                return_dict=True,
            )
            next_logits = out.logits[:, -1, :]

            if do_sample:
                probs = F.softmax(next_logits / safe_temp, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1)
            else:
                next_token = torch.argmax(next_logits, dim=-1, keepdim=True)

            generated_tokens.append(next_token)
            if eos_ids and int(next_token[0, 0].item()) in eos_ids:
                break

            next_embed = self.embed_tokens(next_token)
            cur_ids = torch.cat([cur_ids, next_token], dim=1)
            cur_embeds = torch.cat([cur_embeds, next_embed], dim=1)
            cur_attn = torch.cat(
                [
                    cur_attn,
                    torch.ones(
                        (cur_attn.size(0), 1),
                        dtype=cur_attn.dtype,
                        device=cur_attn.device,
                    ),
                ],
                dim=1,
            )

        if not generated_tokens:
            return start_ids.new_empty((start_ids.size(0), 0))
        return torch.cat(generated_tokens, dim=1)

    @torch.no_grad()
    def generate(
        self,
        input_features: torch.FloatTensor,
        feature_attention_mask: Optional[torch.LongTensor] = None,
        max_new_tokens: int = 128,
        use_baseline: bool = False,
        return_thoughts: bool = False,
        return_stats: bool = False,
        language_hint: Optional[str] = None,
        prompt_text: Optional[str] = None,
        **gen_kwargs: Any,
    ) -> Any:
        """Generate a transcription with or without latent prompting."""
        language_hint = gen_kwargs.pop("language_hint", language_hint)
        prompt_text = gen_kwargs.pop("prompt_text", prompt_text)
        dynamic_halt_threshold = gen_kwargs.pop("dynamic_halt_threshold", getattr(self, "halt_threshold", 0.0))

        B = input_features.size(0)
        assert B == 1, "Evaluation assumes batch size 1."
        device = input_features.device

        if feature_attention_mask is None:
            feature_attention_mask = torch.ones(
                (B, input_features.size(2)), dtype=torch.long, device=device
            )
        audio_lengths = self._get_audio_lengths(input_features, feature_attention_mask)

        pad_token_id = self.processor.tokenizer.pad_token_id
        if pad_token_id is None:
            pad_token_id = self.processor.tokenizer.eos_token_id or 151643

        l_audio = audio_lengths[0]
        audio_toks = build_audio_token_seq(
            self.audio_bos_token_id, self.audio_token_id, self.audio_eos_token_id,
            l_audio, device,
        ).unsqueeze(0)

        system_turn, user_prefix, user_prompt, user_suffix, assistant_prefix = (
            self._build_chat_template_tensors(device, language=language_hint or "English")
        )
        prefix = torch.cat(
            [system_turn, user_prefix, audio_toks.squeeze(0), user_prompt, user_suffix, assistant_prefix],
            dim=0,
        ).unsqueeze(0)

        latent_active = (not use_baseline) and self.use_latent
        soft_prompt_active = (not use_baseline) and self.use_soft_prompt
        interleaved_active = latent_active and self.thought_mode == "interleaved"

        # For prefix mode: prepend NT tokens to the prompt.
        # For interleaved mode: no upfront NT tokens; they are injected during generation.
        front_prompt_active = latent_active or soft_prompt_active
        if front_prompt_active and not interleaved_active:
            nt_toks = torch.tensor([self.nt_token_id] * self.n_latent, device=device).unsqueeze(0)
            start_ids = torch.cat([prefix, nt_toks], dim=1)
        else:
            start_ids = prefix

        attention_mask = (start_ids != pad_token_id).long()
        inputs_embeds = self._get_native_audio_embeds(
            input_ids=start_ids,
            input_features=input_features,
            feature_attention_mask=feature_attention_mask,
            attention_mask=attention_mask,
        )

        states = None
        stats = {}
        if interleaved_active:
            # Interleaved generate: NT tokens injected token-by-token.
            fallback_do_sample = bool(gen_kwargs.get("do_sample", False))
            fallback_temperature = float(gen_kwargs.get("temperature", 1.0))
            gen_ids = self._generate_interleaved(
                base_embeds=inputs_embeds,
                attention_mask=attention_mask,
                max_new_tokens=max_new_tokens,
                do_sample=fallback_do_sample,
                temperature=fallback_temperature,
            )
            if return_thoughts:
                return gen_ids, states
            return gen_ids

        if latent_active:
            initial_state, prefix_past_kv, prefix_attn = self._compute_prefix_state(
                inputs_embeds,
                start_ids,
                strict=True,
                language=language_hint or "English",
            )
            deltas, states, stats = self._compute_continuous_thoughts(
                initial_state,
                prefix_past_key_values=prefix_past_kv,
                prefix_attention_mask=prefix_attn,
                halt_threshold=dynamic_halt_threshold,
            )

            K = deltas.size(1)
            if K == 0:
                # N=0 skip: remove ALL NT tokens, completely reverting to base generation.
                remove_count = self.n_latent
                start_ids = start_ids[:, :-remove_count]
                inputs_embeds = inputs_embeds[:, :-remove_count, :]
                attention_mask = attention_mask[:, :-remove_count]
            elif K < self.n_latent:
                # N > 0 but early halted: Remove the unneeded NT tokens instead of zero-padding them.
                # This prevents the LLM from attending to "blank" un-updated base_nt embeddings.
                remove_count = self.n_latent - K
                start_ids = start_ids[:, :-remove_count]
                inputs_embeds = inputs_embeds[:, :-remove_count, :]
                attention_mask = attention_mask[:, :-remove_count]
                inputs_embeds = self._inject_latent_deltas(
                    inputs_embeds, start_ids, deltas
                )
            else:
                inputs_embeds = self._inject_latent_deltas(
                    inputs_embeds, start_ids, deltas
                )
        elif soft_prompt_active:
            nt_mask = (start_ids == self.nt_token_id)
            nt_counts = nt_mask.sum(dim=1)
            assert (nt_counts == self.n_latent).all(), f"Gen NT count mismatch! {nt_counts}"
            prompt = self.soft_prompt_embed.to(dtype=inputs_embeds.dtype)
            batch_idx, seq_idx = torch.nonzero(nt_mask, as_tuple=True)
            inputs_embeds = inputs_embeds.clone()
            inputs_embeds[batch_idx, seq_idx, :] = prompt.view(-1, self._hidden_size)

        if front_prompt_active:
            is_nt = (start_ids == self.nt_token_id)
            nt_counts = is_nt.sum(dim=1)
            assert (nt_counts <= self.n_latent).all(), f"Gen NT count mismatch! {nt_counts}"

        fallback_eos = gen_kwargs.get("eos_token_id", None)
        fallback_do_sample = bool(gen_kwargs.get("do_sample", False))
        fallback_temperature = float(gen_kwargs.get("temperature", 1.0))

        try:
            if latent_active and K == 0:
                # N=0 skip -> we stripped the NT tokens from start_ids.
                # Do NOT pass inputs_embeds so Qwen generates organically from text + audio ids.
                # Must provide input_features so audio encoder handles the raw audio.
                gen_out = self.thinker.generate(
                    input_ids=start_ids,
                    attention_mask=attention_mask,
                    input_features=input_features,
                    feature_attention_mask=feature_attention_mask,
                    pad_token_id=pad_token_id,
                    max_new_tokens=max_new_tokens,
                    **gen_kwargs,
                )
            else:
                gen_out = self.thinker.generate(
                    inputs_embeds=inputs_embeds,
                    input_ids=start_ids,
                    attention_mask=attention_mask,
                    input_features=None,
                    pad_token_id=pad_token_id,
                    max_new_tokens=max_new_tokens,
                    **gen_kwargs,
                )
            start_len = start_ids.size(1)
            gen_ids = gen_out[:, start_len:]
        except ValueError as e:
            msg = str(e)
            if "inputs_embeds" not in msg:
                raise
            if not self._warned_generate_fallback:
                print(
                    "[warn] thinker.generate does not support inputs_embeds; "
                    "using fallback decoding loop."
                )
                self._warned_generate_fallback = True
            gen_ids = self._generate_with_forward_fallback(
                start_ids=start_ids,
                inputs_embeds=inputs_embeds,
                attention_mask=attention_mask,
                max_new_tokens=max_new_tokens,
                eos_token_id=fallback_eos,
                do_sample=fallback_do_sample,
                temperature=fallback_temperature,
            )
        if return_stats:
            if return_thoughts and latent_active:
                return gen_ids, states, stats
            return gen_ids, stats

        if return_thoughts and latent_active:
            return gen_ids, states
        return gen_ids