qijiv2000

Upload folder using huggingface_hub

46fc068 verified 28 days ago

17.3 kB

"""变体3 hidden2cv 映射层: ST LM hidden 直接映射进 CV3 LLM 文本 embedding 位。
与 bridge (bridge/model.py TextBridge) 的关键区别: 不做 e_cv 蒸馏, 训练信号只有
codec CE (teacher-forcing 穿过冻结 CV3 LLM), 允许表示离开文本嵌入流形。
设计: 零初始化残差 —— out 层零初始化, 初始 e_hat == e (恰在 CV3 文本嵌入流形上,
val_ce 起点应等于冻结 LLM 自身地板 ~2.95), 训练逐步学习 off-manifold 偏移 delta。
结构: LN(h)->proj + LN(e)->proj + 可学习位置 -> pre-norm 双向 Transformer -> 零初始化输出层。
"""
import torch
from torch import nn


class _CausalTransformerBlock(nn.Module):
    """Small causal Transformer block with an explicit KV cache.
    ``nn.TransformerEncoder`` can apply a causal mask for a complete sequence,
    but it does not expose the key/value tensors needed by a long-lived
    streaming session.  This block keeps those tensors in the mapper state and
    therefore has exactly the same computation when a sequence is supplied in
    one call or in several chunks.
    """

    def __init__(self, d_model, n_head, d_ff, dropout):
        super().__init__()
        if d_model % n_head:
            raise ValueError(f'd_model ({d_model}) must be divisible by n_head ({n_head})')
        self.d_model = d_model
        self.n_head = n_head
        self.head_dim = d_model // n_head
        self.norm1 = nn.LayerNorm(d_model)
        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.proj = nn.Linear(d_model, d_model)
        self.attn_drop = nn.Dropout(dropout)
        self.norm2 = nn.LayerNorm(d_model)
        self.ff = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout),
        )

    def _split_heads(self, x):
        # (B,T,D) -> (B,H,T,D/H)
        b, t, _ = x.shape
        return x.view(b, t, self.n_head, self.head_dim).transpose(1, 2)

    def forward(self, x, cache=None, key_padding_mask=None):
        """Run a chunk and return ``(output, new_cache)``.
        ``cache`` is ``(key, value, key_padding_mask)``.  The mask component
        is optional and is only used when callers feed padded chunks.  A
        regular streaming caller can leave it as ``None``.
        """
        b, t, _ = x.shape
        qkv = self.qkv(self.norm1(x)).chunk(3, dim=-1)
        q, k_new, v_new = (self._split_heads(v) for v in qkv)

        past_k = past_v = past_kpm = None
        if cache is not None:
            # Keep accepting a two-tensor cache for callers that construct one
            # themselves; the third item was added for padded training batches.
            past_k, past_v = cache[:2]
            if len(cache) > 2:
                past_kpm = cache[2]
        past_len = 0 if past_k is None else past_k.shape[2]
        k = k_new if past_k is None else torch.cat([past_k, k_new], dim=2)
        v = v_new if past_v is None else torch.cat([past_v, v_new], dim=2)

        # Query positions are [past_len, past_len + t).  A chunk must not
        # attend to a future token from that same chunk.
        total_len = k.shape[2]
        q_pos = torch.arange(past_len, past_len + t, device=x.device)
        k_pos = torch.arange(total_len, device=x.device)
        future = k_pos.unsqueeze(0) > q_pos.unsqueeze(1)  # (T,total_len)

        cur_kpm = None
        if key_padding_mask is not None:
            cur_kpm = key_padding_mask.to(torch.bool)
            if cur_kpm.ndim != 2 or cur_kpm.shape != (b, t):
                raise ValueError(
                    'key_padding_mask must have shape (batch, chunk_length)')
        if past_len:
            # A cache created without a padding mask still represents valid
            # (non-padding) prefix rows.  Materialize those zeros when a later
            # chunk supplies a mask so the mask always spans all keys.
            if past_kpm is None:
                past_kpm = torch.zeros(
                    b, past_len, dtype=torch.bool, device=x.device)
            if cur_kpm is None:
                cur_kpm = torch.zeros(b, t, dtype=torch.bool, device=x.device)
            all_kpm = torch.cat([past_kpm, cur_kpm], dim=1)
        else:
            all_kpm = cur_kpm
        mask = future.unsqueeze(0).expand(b, -1, -1)
        if all_kpm is not None:
            mask = mask | all_kpm.unsqueeze(1)

        scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
        # ``-inf`` yields NaNs for a chunk consisting entirely of padding.  A
        # large finite value keeps that corner case well-defined; its output is
        # masked back to the residual input by the mapper below.
        scores = scores.masked_fill(mask.unsqueeze(1), -1e4)
        attn = torch.softmax(scores, dim=-1)
        attn = self.attn_drop(attn)
        attn = torch.matmul(attn, v).transpose(1, 2).contiguous()
        attn = attn.view(b, t, self.d_model)
        x = x + self.proj(attn)
        x = x + self.ff(self.norm2(x))

        new_kpm = all_kpm
        return x, (k, v, new_kpm)


class Hidden2CVMapper(nn.Module):
    """(h (B,m,2048), e (B,m,896), kpm (B,m)) -> e_hat (B,m,896)"""

    def __init__(self, d_h=2048, d_e=896, n_layer=3, n_head=8, d_ff=2048,
                 dropout=0.1, max_len=512, n_lang=0):
        super().__init__()
        self.d_e = d_e
        self.n_lang = n_lang
        if n_lang > 0:
            # 语种条件 (v3): 零初始化, 不破坏 sanity 地板 (e_hat==e)
            self.lang_emb = nn.Embedding(n_lang, d_e)
            nn.init.zeros_(self.lang_emb.weight)
        self.h_norm = nn.LayerNorm(d_h)
        self.h_proj = nn.Linear(d_h, d_e)
        self.e_norm = nn.LayerNorm(d_e)
        self.e_proj = nn.Linear(d_e, d_e)
        self.pos = nn.Parameter(torch.zeros(1, max_len, d_e))
        nn.init.trunc_normal_(self.pos, std=0.02)
        layer = nn.TransformerEncoderLayer(
            d_e, n_head, d_ff, dropout, activation='gelu',
            batch_first=True, norm_first=True)
        self.encoder = nn.TransformerEncoder(layer, n_layer)
        self.out = nn.Linear(d_e, d_e)
        nn.init.zeros_(self.out.weight)
        nn.init.zeros_(self.out.bias)

    def forward(self, h, e, kpm=None, lang=None):
        """h (B,m,d_h) float, e (B,m,d_e) float, kpm (B,m) bool True=pad,
        lang (B,) long 语种 id (n_lang>0 时必传) -> e_hat (B,m,d_e)"""
        m = h.shape[1]
        x = self.h_proj(self.h_norm(h)) + self.e_proj(self.e_norm(e)) + self.pos[:, :m]
        if self.n_lang > 0:
            assert lang is not None, 'n_lang>0 requires lang ids'
            x = x + self.lang_emb(lang).unsqueeze(1)
        x = self.encoder(x, src_key_padding_mask=kpm)
        return e + self.out(x)


class CausalHidden2CVMapper(nn.Module):
    """Causal, stateful hidden-to-CosyVoice embedding mapper.
    The public input/output contract matches :class:`Hidden2CVMapper`, while
    ``forward_incremental`` and ``step`` keep a per-layer KV cache for a
    persistent streaming session::
        y_offline = mapper(h, e)
        y_a, state = mapper.forward_incremental(h[:, :8], e[:, :8])
        y_b, state = mapper.forward_incremental(h[:, 8:], e[:, 8:], state)
        assert torch.allclose(torch.cat([y_a, y_b], 1), y_offline)
    In residual mode the final projection is zero-initialized, so a new
    mapper is an identity adapter (``e_hat == e``), matching the initialization
    and checkpoint sanity behavior of ``Hidden2CVMapper``.  ``forward`` and the
    incremental path share the same block implementation, so there is no
    train/inference-only causal mask discrepancy.  For exact numerical parity
    between full and chunked calls, use ``eval()`` or ``dropout=0``.
    ``state`` is an opaque dictionary containing ``position`` and one
    ``(key, value, key_padding_mask)`` cache tuple per layer.  Callers should
    pass it back unchanged to the next invocation; ``reset_state`` starts a
    fresh sequence.
    ``input_mode='h_only'`` removes the native e projections and residual. It
    uses only h, language and position to predict the CV embedding; passing e
    then supplies shape/dtype compatibility but never its values.
    """

    def __init__(self, d_h=2048, d_e=896, n_layer=3, n_head=8, d_ff=2048,
                 dropout=0.1, max_len=512, n_lang=0, input_mode='residual'):
        super().__init__()
        if input_mode not in ('residual', 'h_only'):
            raise ValueError("input_mode must be 'residual' or 'h_only'")
        self.input_mode = input_mode
        self.d_h = d_h
        self.d_e = d_e
        self.n_layer = n_layer
        self.n_head = n_head
        self.d_ff = d_ff
        self.dropout = dropout
        self.max_len = max_len
        self.n_lang = n_lang
        if n_lang > 0:
            # Keep language conditioning zero-initialized so adding it cannot
            # alter the e_hat == e initialization sanity check.
            self.lang_emb = nn.Embedding(n_lang, d_e)
            nn.init.zeros_(self.lang_emb.weight)
        self.h_norm = nn.LayerNorm(d_h)
        self.h_proj = nn.Linear(d_h, d_e)
        if input_mode == 'residual':
            self.e_norm = nn.LayerNorm(d_e)
            self.e_proj = nn.Linear(d_e, d_e)
        self.pos = nn.Parameter(torch.zeros(1, max_len, d_e))
        nn.init.trunc_normal_(self.pos, std=0.02)
        self.layers = nn.ModuleList([
            _CausalTransformerBlock(d_e, n_head, d_ff, dropout)
            for _ in range(n_layer)
        ])
        self.out = nn.Linear(d_e, d_e)
        if input_mode == 'residual':
            nn.init.zeros_(self.out.weight)
        else:
            # h-only arms must not have an embedding lookup/residual shortcut.
            nn.init.xavier_uniform_(self.out.weight)
        nn.init.zeros_(self.out.bias)

    def _position_slice(self, start, length, device, dtype):
        """Return deterministic position features, including beyond max_len.
        Existing checkpoints are trained with ``max_len`` learned positions.
        For a longer streaming sequence, learned positions are used where
        available and zeros are appended.  The recurrent KV context still
        distinguishes all later tokens and, crucially, 60-second inputs do not
        fail merely because they exceed the training position table.
        """
        if length == 0:
            return self.pos.new_zeros((1, 0, self.d_e), device=device, dtype=dtype)
        end = start + length
        if start >= self.max_len:
            return self.pos.new_zeros((1, length, self.d_e), device=device, dtype=dtype)
        learned_end = min(end, self.max_len)
        learned = self.pos[:, start:learned_end].to(device=device, dtype=dtype)
        if learned_end == end:
            return learned
        pad = self.pos.new_zeros(
            (1, end - learned_end, self.d_e), device=device, dtype=dtype)
        return torch.cat([learned, pad], dim=1)

    def reset_state(self, batch_size=1, device=None, dtype=None):
        """Create an empty streaming state for ``batch_size`` sequences."""
        if device is None:
            device = self.pos.device
        if dtype is None:
            dtype = self.pos.dtype
        # KV tensors are created by the first call, so dtype is retained here
        # only for API symmetry and future callers that inspect the state.
        del dtype
        return {
            'position': 0,
            'kv': [None for _ in self.layers],
            'batch_size': int(batch_size),
            'device': torch.device(device),
        }

    # Alias used by a few streaming wrappers.
    init_state = reset_state

    def _validate_inputs(self, h, e):
        if h.ndim == 2:
            h = h.unsqueeze(1)
        if e is None:
            if self.input_mode != 'h_only':
                raise ValueError('residual mapper requires native embeddings e')
            e = h.new_zeros((*h.shape[:2], self.d_e))
        if e.ndim == 2:
            e = e.unsqueeze(1)
        if h.ndim != 3 or e.ndim != 3:
            raise ValueError('h and e must have shape (B,T,D) or (B,D)')
        if h.shape[:2] != e.shape[:2]:
            raise ValueError('h and e must have the same batch and sequence dimensions')
        if h.shape[-1] != self.d_h:
            raise ValueError(f'h last dimension must be {self.d_h}, got {h.shape[-1]}')
        if e.shape[-1] != self.d_e:
            raise ValueError(f'e last dimension must be {self.d_e}, got {e.shape[-1]}')
        return h, e

    def _prepare_state(self, state, batch_size, device):
        if state is None:
            return self.reset_state(batch_size, device=device)
        if not isinstance(state, dict):
            raise TypeError('state must be a state dictionary returned by the mapper')
        kv = state.get('kv')
        if kv is None or len(kv) != len(self.layers):
            raise ValueError('state has an invalid per-layer KV cache')
        if state.get('batch_size', batch_size) != batch_size:
            raise ValueError('state batch size does not match h/e')
        if state.get('position', 0) < 0:
            raise ValueError('state position must be non-negative')
        # Do not mutate the caller's dictionary.  The cache tensors themselves
        # are immutable references until a layer concatenates a new chunk.
        return {
            'position': int(state.get('position', 0)),
            'kv': list(kv),
            'batch_size': batch_size,
            'device': torch.device(device),
        }

    def forward_incremental(self, h, e=None, state=None, lang=None, kpm=None):
        """Map one sequence/chunk and return ``(e_hat, new_state)``.
        Args:
            h: ST hidden states, ``(B,T,d_h)`` or one-token ``(B,d_h)``.
            e: CV3 text embeddings, ``(B,T,d_e)`` or ``(B,d_e)``.
            state: State returned by an earlier call, or ``None`` to reset.
            lang: Optional ``(B,)`` language ids when ``n_lang > 0``.
            kpm: Optional ``(B,T)`` padding mask (``True`` means padding).
        """
        h, e = self._validate_inputs(h, e)
        b, t, _ = h.shape
        if kpm is not None:
            if kpm.shape != (b, t):
                raise ValueError('kpm must have shape (batch, sequence_length)')
            kpm = kpm.to(dtype=torch.bool, device=h.device)
        if self.n_lang > 0:
            if lang is None:
                raise AssertionError('n_lang>0 requires lang ids')
            if lang.ndim == 0:
                lang = lang.reshape(1).expand(b)
            if lang.shape != (b,):
                raise ValueError('lang must have shape (batch,)')
            lang = lang.to(device=h.device, dtype=torch.long)

        st = self._prepare_state(state, b, h.device)
        if t == 0:
            return e, st
        pos = self._position_slice(st['position'], t, h.device, e.dtype)
        x = self.h_proj(self.h_norm(h)) + pos
        if self.input_mode == 'residual':
            x = x + self.e_proj(self.e_norm(e))
        if self.n_lang > 0:
            x = x + self.lang_emb(lang).unsqueeze(1).to(dtype=x.dtype)

        new_kv = []
        for layer, cache in zip(self.layers, st['kv']):
            x, cache = layer(x, cache=cache, key_padding_mask=kpm)
            new_kv.append(cache)

        delta = self.out(x)
        if kpm is not None:
            # A padded position must not update either its output or the
            # semantic prefix represented by the state.
            delta = delta.masked_fill(kpm.unsqueeze(-1), 0)
        e_hat = e + delta if self.input_mode == 'residual' else delta
        st['position'] += t
        st['kv'] = new_kv
        return e_hat, st

    def step(self, h, e=None, state=None, lang=None, kpm=None):
        """Process one token and return ``(e_hat_t, new_state)``.
        ``h``/``e`` may be either ``(B,D)`` or ``(B,1,D)``.  This is a thin
        wrapper around ``forward_incremental`` so both entry points have
        identical cache semantics.
        """
        if h.ndim == 3 and h.shape[1] != 1:
            raise ValueError('step requires exactly one token')
        h_hat, st = self.forward_incremental(h, e, state=state, lang=lang, kpm=kpm)
        return h_hat[:, 0], st

    def forward(self, h, e=None, kpm=None, lang=None, chunk_sizes=None):
        """Offline causal forward; equivalent to one incremental call."""
        if chunk_sizes is not None:
            h, e = self._validate_inputs(h, e)
            if any(size <= 0 for size in chunk_sizes) or sum(chunk_sizes) != h.shape[1]:
                raise ValueError('chunk_sizes must be positive and cover the complete sequence')
            state, outputs, start = None, [], 0
            for size in chunk_sizes:
                end = start + size
                out, state = self.forward_incremental(
                    h[:, start:end], e[:, start:end], state=state, lang=lang,
                    kpm=None if kpm is None else kpm[:, start:end])
                outputs.append(out)
                start = end
            return torch.cat(outputs, dim=1) if outputs else e
        e_hat, _ = self.forward_incremental(h, e, state=None, lang=lang, kpm=kpm)
        return e_hat


__all__ = ['Hidden2CVMapper', 'CausalHidden2CVMapper']

	"""变体3 hidden2cv 映射层: ST LM hidden 直接映射进 CV3 LLM 文本 embedding 位。

	与 bridge (bridge/model.py TextBridge) 的关键区别: 不做 e_cv 蒸馏, 训练信号只有
	codec CE (teacher-forcing 穿过冻结 CV3 LLM), 允许表示离开文本嵌入流形。

	设计: 零初始化残差 —— out 层零初始化, 初始 e_hat == e (恰在 CV3 文本嵌入流形上,
	val_ce 起点应等于冻结 LLM 自身地板 ~2.95), 训练逐步学习 off-manifold 偏移 delta。
	结构: LN(h)->proj + LN(e)->proj + 可学习位置 -> pre-norm 双向 Transformer -> 零初始化输出层。
	"""
	import torch
	from torch import nn


	class _CausalTransformerBlock(nn.Module):
	"""Small causal Transformer block with an explicit KV cache.

	``nn.TransformerEncoder`` can apply a causal mask for a complete sequence,
	but it does not expose the key/value tensors needed by a long-lived
	streaming session. This block keeps those tensors in the mapper state and
	therefore has exactly the same computation when a sequence is supplied in
	one call or in several chunks.
	"""

	def __init__(self, d_model, n_head, d_ff, dropout):
	super().__init__()
	if d_model % n_head:
	raise ValueError(f'd_model ({d_model}) must be divisible by n_head ({n_head})')
	self.d_model = d_model
	self.n_head = n_head
	self.head_dim = d_model // n_head
	self.norm1 = nn.LayerNorm(d_model)
	self.qkv = nn.Linear(d_model, 3 * d_model)
	self.proj = nn.Linear(d_model, d_model)
	self.attn_drop = nn.Dropout(dropout)
	self.norm2 = nn.LayerNorm(d_model)
	self.ff = nn.Sequential(
	nn.Linear(d_model, d_ff),
	nn.GELU(),
	nn.Dropout(dropout),
	nn.Linear(d_ff, d_model),
	nn.Dropout(dropout),
	)

	def _split_heads(self, x):
	# (B,T,D) -> (B,H,T,D/H)
	b, t, _ = x.shape
	return x.view(b, t, self.n_head, self.head_dim).transpose(1, 2)

	def forward(self, x, cache=None, key_padding_mask=None):
	"""Run a chunk and return ``(output, new_cache)``.

	``cache`` is ``(key, value, key_padding_mask)``. The mask component
	is optional and is only used when callers feed padded chunks. A
	regular streaming caller can leave it as ``None``.
	"""
	b, t, _ = x.shape
	qkv = self.qkv(self.norm1(x)).chunk(3, dim=-1)
	q, k_new, v_new = (self._split_heads(v) for v in qkv)

	past_k = past_v = past_kpm = None
	if cache is not None:
	# Keep accepting a two-tensor cache for callers that construct one
	# themselves; the third item was added for padded training batches.
	past_k, past_v = cache[:2]
	if len(cache) > 2:
	past_kpm = cache[2]
	past_len = 0 if past_k is None else past_k.shape[2]
	k = k_new if past_k is None else torch.cat([past_k, k_new], dim=2)
	v = v_new if past_v is None else torch.cat([past_v, v_new], dim=2)

	# Query positions are [past_len, past_len + t). A chunk must not
	# attend to a future token from that same chunk.
	total_len = k.shape[2]
	q_pos = torch.arange(past_len, past_len + t, device=x.device)
	k_pos = torch.arange(total_len, device=x.device)
	future = k_pos.unsqueeze(0) > q_pos.unsqueeze(1) # (T,total_len)

	cur_kpm = None
	if key_padding_mask is not None:
	cur_kpm = key_padding_mask.to(torch.bool)
	if cur_kpm.ndim != 2 or cur_kpm.shape != (b, t):
	raise ValueError(
	'key_padding_mask must have shape (batch, chunk_length)')
	if past_len:
	# A cache created without a padding mask still represents valid
	# (non-padding) prefix rows. Materialize those zeros when a later
	# chunk supplies a mask so the mask always spans all keys.
	if past_kpm is None:
	past_kpm = torch.zeros(
	b, past_len, dtype=torch.bool, device=x.device)
	if cur_kpm is None:
	cur_kpm = torch.zeros(b, t, dtype=torch.bool, device=x.device)
	all_kpm = torch.cat([past_kpm, cur_kpm], dim=1)
	else:
	all_kpm = cur_kpm
	mask = future.unsqueeze(0).expand(b, -1, -1)
	if all_kpm is not None:
	mask = mask \| all_kpm.unsqueeze(1)

	scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
	# ``-inf`` yields NaNs for a chunk consisting entirely of padding. A
	# large finite value keeps that corner case well-defined; its output is
	# masked back to the residual input by the mapper below.
	scores = scores.masked_fill(mask.unsqueeze(1), -1e4)
	attn = torch.softmax(scores, dim=-1)
	attn = self.attn_drop(attn)
	attn = torch.matmul(attn, v).transpose(1, 2).contiguous()
	attn = attn.view(b, t, self.d_model)
	x = x + self.proj(attn)
	x = x + self.ff(self.norm2(x))

	new_kpm = all_kpm
	return x, (k, v, new_kpm)


	class Hidden2CVMapper(nn.Module):
	"""(h (B,m,2048), e (B,m,896), kpm (B,m)) -> e_hat (B,m,896)"""

	def __init__(self, d_h=2048, d_e=896, n_layer=3, n_head=8, d_ff=2048,
	dropout=0.1, max_len=512, n_lang=0):
	super().__init__()
	self.d_e = d_e
	self.n_lang = n_lang
	if n_lang > 0:
	# 语种条件 (v3): 零初始化, 不破坏 sanity 地板 (e_hat==e)
	self.lang_emb = nn.Embedding(n_lang, d_e)
	nn.init.zeros_(self.lang_emb.weight)
	self.h_norm = nn.LayerNorm(d_h)
	self.h_proj = nn.Linear(d_h, d_e)
	self.e_norm = nn.LayerNorm(d_e)
	self.e_proj = nn.Linear(d_e, d_e)
	self.pos = nn.Parameter(torch.zeros(1, max_len, d_e))
	nn.init.trunc_normal_(self.pos, std=0.02)
	layer = nn.TransformerEncoderLayer(
	d_e, n_head, d_ff, dropout, activation='gelu',
	batch_first=True, norm_first=True)
	self.encoder = nn.TransformerEncoder(layer, n_layer)
	self.out = nn.Linear(d_e, d_e)
	nn.init.zeros_(self.out.weight)
	nn.init.zeros_(self.out.bias)

	def forward(self, h, e, kpm=None, lang=None):
	"""h (B,m,d_h) float, e (B,m,d_e) float, kpm (B,m) bool True=pad,
	lang (B,) long 语种 id (n_lang>0 时必传) -> e_hat (B,m,d_e)"""
	m = h.shape[1]
	x = self.h_proj(self.h_norm(h)) + self.e_proj(self.e_norm(e)) + self.pos[:, :m]
	if self.n_lang > 0:
	assert lang is not None, 'n_lang>0 requires lang ids'
	x = x + self.lang_emb(lang).unsqueeze(1)
	x = self.encoder(x, src_key_padding_mask=kpm)
	return e + self.out(x)


	class CausalHidden2CVMapper(nn.Module):
	"""Causal, stateful hidden-to-CosyVoice embedding mapper.

	The public input/output contract matches :class:`Hidden2CVMapper`, while
	``forward_incremental`` and ``step`` keep a per-layer KV cache for a
	persistent streaming session::

	y_offline = mapper(h, e)
	y_a, state = mapper.forward_incremental(h[:, :8], e[:, :8])
	y_b, state = mapper.forward_incremental(h[:, 8:], e[:, 8:], state)
	assert torch.allclose(torch.cat([y_a, y_b], 1), y_offline)

	In residual mode the final projection is zero-initialized, so a new
	mapper is an identity adapter (``e_hat == e``), matching the initialization
	and checkpoint sanity behavior of ``Hidden2CVMapper``. ``forward`` and the
	incremental path share the same block implementation, so there is no
	train/inference-only causal mask discrepancy. For exact numerical parity
	between full and chunked calls, use ``eval()`` or ``dropout=0``.

	``state`` is an opaque dictionary containing ``position`` and one
	``(key, value, key_padding_mask)`` cache tuple per layer. Callers should
	pass it back unchanged to the next invocation; ``reset_state`` starts a
	fresh sequence.

	``input_mode='h_only'`` removes the native e projections and residual. It
	uses only h, language and position to predict the CV embedding; passing e
	then supplies shape/dtype compatibility but never its values.
	"""

	def __init__(self, d_h=2048, d_e=896, n_layer=3, n_head=8, d_ff=2048,
	dropout=0.1, max_len=512, n_lang=0, input_mode='residual'):
	super().__init__()
	if input_mode not in ('residual', 'h_only'):
	raise ValueError("input_mode must be 'residual' or 'h_only'")
	self.input_mode = input_mode
	self.d_h = d_h
	self.d_e = d_e
	self.n_layer = n_layer
	self.n_head = n_head
	self.d_ff = d_ff
	self.dropout = dropout
	self.max_len = max_len
	self.n_lang = n_lang
	if n_lang > 0:
	# Keep language conditioning zero-initialized so adding it cannot
	# alter the e_hat == e initialization sanity check.
	self.lang_emb = nn.Embedding(n_lang, d_e)
	nn.init.zeros_(self.lang_emb.weight)
	self.h_norm = nn.LayerNorm(d_h)
	self.h_proj = nn.Linear(d_h, d_e)
	if input_mode == 'residual':
	self.e_norm = nn.LayerNorm(d_e)
	self.e_proj = nn.Linear(d_e, d_e)
	self.pos = nn.Parameter(torch.zeros(1, max_len, d_e))
	nn.init.trunc_normal_(self.pos, std=0.02)
	self.layers = nn.ModuleList([
	_CausalTransformerBlock(d_e, n_head, d_ff, dropout)
	for _ in range(n_layer)
	])
	self.out = nn.Linear(d_e, d_e)
	if input_mode == 'residual':
	nn.init.zeros_(self.out.weight)
	else:
	# h-only arms must not have an embedding lookup/residual shortcut.
	nn.init.xavier_uniform_(self.out.weight)
	nn.init.zeros_(self.out.bias)

	def _position_slice(self, start, length, device, dtype):
	"""Return deterministic position features, including beyond max_len.

	Existing checkpoints are trained with ``max_len`` learned positions.
	For a longer streaming sequence, learned positions are used where
	available and zeros are appended. The recurrent KV context still
	distinguishes all later tokens and, crucially, 60-second inputs do not
	fail merely because they exceed the training position table.
	"""
	if length == 0:
	return self.pos.new_zeros((1, 0, self.d_e), device=device, dtype=dtype)
	end = start + length
	if start >= self.max_len:
	return self.pos.new_zeros((1, length, self.d_e), device=device, dtype=dtype)
	learned_end = min(end, self.max_len)
	learned = self.pos[:, start:learned_end].to(device=device, dtype=dtype)
	if learned_end == end:
	return learned
	pad = self.pos.new_zeros(
	(1, end - learned_end, self.d_e), device=device, dtype=dtype)
	return torch.cat([learned, pad], dim=1)

	def reset_state(self, batch_size=1, device=None, dtype=None):
	"""Create an empty streaming state for ``batch_size`` sequences."""
	if device is None:
	device = self.pos.device
	if dtype is None:
	dtype = self.pos.dtype
	# KV tensors are created by the first call, so dtype is retained here
	# only for API symmetry and future callers that inspect the state.
	del dtype
	return {
	'position': 0,
	'kv': [None for _ in self.layers],
	'batch_size': int(batch_size),
	'device': torch.device(device),
	}

	# Alias used by a few streaming wrappers.
	init_state = reset_state

	def _validate_inputs(self, h, e):
	if h.ndim == 2:
	h = h.unsqueeze(1)
	if e is None:
	if self.input_mode != 'h_only':
	raise ValueError('residual mapper requires native embeddings e')
	e = h.new_zeros((*h.shape[:2], self.d_e))
	if e.ndim == 2:
	e = e.unsqueeze(1)
	if h.ndim != 3 or e.ndim != 3:
	raise ValueError('h and e must have shape (B,T,D) or (B,D)')
	if h.shape[:2] != e.shape[:2]:
	raise ValueError('h and e must have the same batch and sequence dimensions')
	if h.shape[-1] != self.d_h:
	raise ValueError(f'h last dimension must be {self.d_h}, got {h.shape[-1]}')
	if e.shape[-1] != self.d_e:
	raise ValueError(f'e last dimension must be {self.d_e}, got {e.shape[-1]}')
	return h, e

	def _prepare_state(self, state, batch_size, device):
	if state is None:
	return self.reset_state(batch_size, device=device)
	if not isinstance(state, dict):
	raise TypeError('state must be a state dictionary returned by the mapper')
	kv = state.get('kv')
	if kv is None or len(kv) != len(self.layers):
	raise ValueError('state has an invalid per-layer KV cache')
	if state.get('batch_size', batch_size) != batch_size:
	raise ValueError('state batch size does not match h/e')
	if state.get('position', 0) < 0:
	raise ValueError('state position must be non-negative')
	# Do not mutate the caller's dictionary. The cache tensors themselves
	# are immutable references until a layer concatenates a new chunk.
	return {
	'position': int(state.get('position', 0)),
	'kv': list(kv),
	'batch_size': batch_size,
	'device': torch.device(device),
	}

	def forward_incremental(self, h, e=None, state=None, lang=None, kpm=None):
	"""Map one sequence/chunk and return ``(e_hat, new_state)``.

	Args:
	h: ST hidden states, ``(B,T,d_h)`` or one-token ``(B,d_h)``.
	e: CV3 text embeddings, ``(B,T,d_e)`` or ``(B,d_e)``.
	state: State returned by an earlier call, or ``None`` to reset.
	lang: Optional ``(B,)`` language ids when ``n_lang > 0``.
	kpm: Optional ``(B,T)`` padding mask (``True`` means padding).
	"""
	h, e = self._validate_inputs(h, e)
	b, t, _ = h.shape
	if kpm is not None:
	if kpm.shape != (b, t):
	raise ValueError('kpm must have shape (batch, sequence_length)')
	kpm = kpm.to(dtype=torch.bool, device=h.device)
	if self.n_lang > 0:
	if lang is None:
	raise AssertionError('n_lang>0 requires lang ids')
	if lang.ndim == 0:
	lang = lang.reshape(1).expand(b)
	if lang.shape != (b,):
	raise ValueError('lang must have shape (batch,)')
	lang = lang.to(device=h.device, dtype=torch.long)

	st = self._prepare_state(state, b, h.device)
	if t == 0:
	return e, st
	pos = self._position_slice(st['position'], t, h.device, e.dtype)
	x = self.h_proj(self.h_norm(h)) + pos
	if self.input_mode == 'residual':
	x = x + self.e_proj(self.e_norm(e))
	if self.n_lang > 0:
	x = x + self.lang_emb(lang).unsqueeze(1).to(dtype=x.dtype)

	new_kv = []
	for layer, cache in zip(self.layers, st['kv']):
	x, cache = layer(x, cache=cache, key_padding_mask=kpm)
	new_kv.append(cache)

	delta = self.out(x)
	if kpm is not None:
	# A padded position must not update either its output or the
	# semantic prefix represented by the state.
	delta = delta.masked_fill(kpm.unsqueeze(-1), 0)
	e_hat = e + delta if self.input_mode == 'residual' else delta
	st['position'] += t
	st['kv'] = new_kv
	return e_hat, st

	def step(self, h, e=None, state=None, lang=None, kpm=None):
	"""Process one token and return ``(e_hat_t, new_state)``.

	``h``/``e`` may be either ``(B,D)`` or ``(B,1,D)``. This is a thin
	wrapper around ``forward_incremental`` so both entry points have
	identical cache semantics.
	"""
	if h.ndim == 3 and h.shape[1] != 1:
	raise ValueError('step requires exactly one token')
	h_hat, st = self.forward_incremental(h, e, state=state, lang=lang, kpm=kpm)
	return h_hat[:, 0], st

	def forward(self, h, e=None, kpm=None, lang=None, chunk_sizes=None):
	"""Offline causal forward; equivalent to one incremental call."""
	if chunk_sizes is not None:
	h, e = self._validate_inputs(h, e)
	if any(size <= 0 for size in chunk_sizes) or sum(chunk_sizes) != h.shape[1]:
	raise ValueError('chunk_sizes must be positive and cover the complete sequence')
	state, outputs, start = None, [], 0
	for size in chunk_sizes:
	end = start + size
	out, state = self.forward_incremental(
	h[:, start:end], e[:, start:end], state=state, lang=lang,
	kpm=None if kpm is None else kpm[:, start:end])
	outputs.append(out)
	start = end
	return torch.cat(outputs, dim=1) if outputs else e
	e_hat, _ = self.forward_incremental(h, e, state=None, lang=lang, kpm=kpm)
	return e_hat


	__all__ = ['Hidden2CVMapper', 'CausalHidden2CVMapper']