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
|
from funasr.register import tables
|
import base64
|
import gzip
|
from dataclasses import dataclass
|
from typing import Dict, Iterable, Optional
|
|
import numpy as np
|
import torch
|
import torch.nn.functional as F
|
from torch import Tensor, nn
|
|
|
class LayerNorm(nn.LayerNorm):
|
def forward(self, x: Tensor) -> Tensor:
|
return super().forward(x.float()).type(x.dtype)
|
|
|
class Linear(nn.Linear):
|
def forward(self, x: Tensor) -> Tensor:
|
return F.linear(
|
x,
|
self.weight.to(x.dtype),
|
None if self.bias is None else self.bias.to(x.dtype),
|
)
|
|
|
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
|
|
|
class MultiHeadAttention(nn.Module):
|
def __init__(self, n_state: int, n_head: int):
|
super().__init__()
|
self.n_head = n_head
|
self.query = Linear(n_state, n_state)
|
self.key = Linear(n_state, n_state, bias=False)
|
self.value = Linear(n_state, n_state)
|
self.out = Linear(n_state, n_state)
|
|
def forward(
|
self,
|
x: Tensor,
|
xa: Optional[Tensor] = None,
|
mask: Optional[Tensor] = None,
|
kv_cache: Optional[dict] = None,
|
**kwargs,
|
):
|
is_pad_mask = kwargs.get("is_pad_mask", False)
|
|
q = self.query(x)
|
|
if kv_cache is None or xa is None or self.key not in kv_cache:
|
# hooks, if installed (i.e. kv_cache is not None), will prepend the cached kv tensors;
|
# otherwise, perform key/value projections for self- or cross-attention as usual.
|
k = self.key(x if xa is None else xa)
|
v = self.value(x if xa is None else xa)
|
else:
|
# for cross-attention, calculate keys and values once and reuse in subsequent calls.
|
k = kv_cache[self.key]
|
v = kv_cache[self.value]
|
|
wv, qk = self.qkv_attention(q, k, v, mask, is_pad_mask=is_pad_mask)
|
return self.out(wv), qk
|
|
def qkv_attention(
|
self,
|
q: Tensor,
|
k: Tensor,
|
v: Tensor,
|
mask: Optional[Tensor] = None,
|
**kwargs,
|
):
|
is_pad_mask = kwargs.get("is_pad_mask", False)
|
n_batch, n_ctx, n_state = q.shape
|
scale = (n_state // self.n_head) ** -0.25
|
q = q.view(*q.shape[:2], self.n_head, -1).permute(0, 2, 1, 3) * scale
|
k = k.view(*k.shape[:2], self.n_head, -1).permute(0, 2, 3, 1) * scale
|
v = v.view(*v.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
|
|
qk = q @ k
|
if mask is not None:
|
if not is_pad_mask:
|
qk = qk + mask[:n_ctx, :n_ctx]
|
else:
|
mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2)
|
min_value = float(np.finfo(torch.tensor(0, dtype=qk.dtype).numpy().dtype).min)
|
qk = qk.masked_fill(mask, min_value)
|
|
qk = qk.float()
|
|
w = F.softmax(qk, dim=-1).to(q.dtype)
|
if mask is not None and is_pad_mask:
|
w = w.masked_fill(mask, 0.0)
|
return (w @ v).permute(0, 2, 1, 3).flatten(start_dim=2), qk.detach()
|
|
|
from omegaconf import OmegaConf
|
|
|
class ResidualAttentionBlockRWKV(nn.Module):
|
def __init__(
|
self, n_state: int, n_head: int, cross_attention: bool = False, layer_id=0, **kwargs
|
):
|
super().__init__()
|
|
rwkv_cfg = kwargs.get("rwkv_cfg", {})
|
args = OmegaConf.create(rwkv_cfg)
|
if args.get("version", "v4") == "v4":
|
from funasr.models.sense_voice.rwkv_v4 import RWKVLayer
|
from funasr.models.sense_voice.rwkv_v4 import RWKV_TimeMix as RWKV_Tmix
|
elif args.get("version", "v5") == "v5":
|
from funasr.models.sense_voice.rwkv_v5 import RWKVLayer
|
from funasr.models.sense_voice.rwkv_v5 import RWKV_Tmix_x052 as RWKV_Tmix
|
else:
|
from funasr.models.sense_voice.rwkv_v6 import RWKVLayer
|
from funasr.models.sense_voice.rwkv_v6 import RWKV_Tmix_x060 as RWKV_Tmix
|
# self.att = RWKVLayer(args=args, layer_id=layer_id)
|
self.att = RWKV_Tmix(args, layer_id=layer_id)
|
|
if args.get("init_rwkv", True):
|
print("init_rwkv")
|
nn.init.orthogonal_(self.att.receptance.weight, gain=1)
|
nn.init.orthogonal_(self.att.key.weight, gain=0.1)
|
nn.init.orthogonal_(self.att.value.weight, gain=1)
|
nn.init.orthogonal_(self.att.gate.weight, gain=0.1)
|
nn.init.zeros_(self.att.output.weight)
|
|
self.ln0 = None
|
if layer_id == 0 and not args.get("ln0", True):
|
self.ln0 = LayerNorm(args.n_embd)
|
if args.get("init_rwkv", True):
|
print("init_rwkv")
|
layer_id = 0
|
scale = ((1 + layer_id) / args.get("n_layer")) ** 0.7
|
nn.init.constant_(self.ln0.weight, scale)
|
|
self.layer_id = layer_id
|
self.args = args
|
|
self.ln1 = None
|
if not args.get("ln1", True):
|
self.ln1 = LayerNorm(args.n_embd)
|
# init
|
if args.get("init_rwkv", True):
|
print("init_rwkv")
|
scale = ((1 + layer_id) / args.get("n_layer")) ** 0.7
|
nn.init.constant_(self.ln1.weight, scale)
|
|
if args.get("datatype", "bf16") == "bf16":
|
self.att.to(torch.bfloat16)
|
# if self.ln1 is not None:
|
# self.ln1.to(torch.bfloat16)
|
|
self.cross_attn = MultiHeadAttention(n_state, n_head) if cross_attention else None
|
self.cross_attn_ln = LayerNorm(n_state) if cross_attention else None
|
|
n_mlp = n_state * 4
|
self.mlp = nn.Sequential(Linear(n_state, n_mlp), nn.GELU(), Linear(n_mlp, n_state))
|
self.mlp_ln = LayerNorm(n_state)
|
|
def forward(
|
self,
|
x: Tensor,
|
xa: Optional[Tensor] = None,
|
mask: Optional[Tensor] = None,
|
kv_cache: Optional[dict] = None,
|
**kwargs,
|
):
|
is_pad_mask = kwargs.get("is_pad_mask", False)
|
is_pad_memory_mask = kwargs.get("is_pad_memory_mask", False)
|
|
if self.layer_id == 0 and self.ln0 is not None:
|
x = self.ln0(x)
|
|
if self.args.get("datatype", "bf16") == "bf16":
|
x = x.bfloat16()
|
if self.ln1 is None:
|
x = x + self.att(x, mask=mask, kv_cache=kv_cache, is_pad_mask=is_pad_mask)[0]
|
else:
|
x = x + self.att(self.ln1(x), mask=mask, kv_cache=kv_cache, is_pad_mask=is_pad_mask)[0]
|
if self.args.get("datatype", "bf16") == "bf16":
|
x = x.to(torch.float32)
|
|
if self.cross_attn:
|
x = (
|
x
|
+ self.cross_attn(
|
self.cross_attn_ln(x), xa, kv_cache=kv_cache, is_pad_mask=is_pad_memory_mask
|
)[0]
|
)
|
x = x + self.mlp(self.mlp_ln(x))
|
|
return x
|
|
|
@tables.register("decoder_classes", "SenseVoiceDecoder")
|
class SenseVoiceDecoder(nn.Module):
|
def __init__(self, n_vocab: int, n_ctx: int, n_state: int, n_head: int, n_layer: int, **kwargs):
|
super().__init__()
|
|
self.token_embedding = nn.Embedding(n_vocab, n_state)
|
self.positional_embedding = nn.Parameter(torch.empty(n_ctx, n_state))
|
|
self.blocks = nn.ModuleList(
|
[
|
ResidualAttentionBlockRWKV(
|
n_state, n_head, cross_attention=True, layer_id=i, **kwargs
|
)
|
for i in range(n_layer)
|
]
|
)
|
self.ln = LayerNorm(n_state)
|
|
mask = torch.empty(n_ctx, n_ctx).fill_(-np.inf).triu_(1)
|
self.register_buffer("mask", mask, persistent=False)
|
|
self.use_padmask = kwargs.get("use_padmask", True)
|
|
def 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
|
