import copy
|
from typing import Optional, Tuple, Union
|
|
import torch
|
import torch.nn as nn
|
import torch.nn.functional as F
|
from funasr.models.transformer.utils.nets_utils import make_pad_mask
|
|
def sense_voice_decode_forward(
|
self,
|
x: torch.Tensor,
|
xa: torch.Tensor,
|
kv_cache: Optional[dict] = None,
|
**kwargs,
|
):
|
"""Forward decoder.
|
|
Args:
|
hs_pad: encoded memory, float32 (batch, maxlen_in, feat)
|
hlens: (batch)
|
ys_in_pad:
|
input token ids, int64 (batch, maxlen_out)
|
if input_layer == "embed"
|
input tensor (batch, maxlen_out, #mels) in the other cases
|
ys_in_lens: (batch)
|
Returns:
|
(tuple): tuple containing:
|
|
x: decoded token score before softmax (batch, maxlen_out, token)
|
if use_output_layer is True,
|
olens: (batch, )
|
"""
|
# import pdb;pdb.set_trace()
|
use_padmask = self.use_padmask
|
hlens = kwargs.get("hlens", None)
|
|
ys_in_lens = kwargs.get("ys_in_lens", None)
|
|
offset = next(iter(kv_cache.values())).shape[1] if kv_cache else 0
|
tgt, memory = x, xa
|
tgt[tgt==-1] = 0
|
tgt = (
|
self.token_embedding(tgt)
|
+ self.positional_embedding[offset : offset + tgt.size(1)]
|
)
|
# tgt = self.dropout(tgt)
|
|
x = tgt.to(memory.dtype)
|
|
if use_padmask and hlens is not None:
|
memory_mask = (~make_pad_mask(hlens)[:, None, :]).to(memory.device)
|
else:
|
memory_mask = None
|
|
for layer, block in enumerate(self.blocks):
|
x = block(x, memory, mask=self.mask, memory_mask=memory_mask, is_pad_mask=False, is_pad_memory_mask=True)
|
|
|
x = self.ln(x)
|
x = (
|
x @ torch.transpose(self.token_embedding.weight.to(x.dtype), 0, 1)
|
).float()
|
|
|
return x
|
|
