first commit for takway.ai
This commit is contained in:
IrvingGao committed
commit
d7a6de771c
3256 files changed
+795173
No files matched your search
Whitespace-only changes.
@@ -0,0 +1,978 @@
|
||||
"""Beam search module."""
|
||||
|
||||
from itertools import chain
|
||||
import logging
|
||||
from typing import Any
|
||||
from typing import Dict
|
||||
from typing import List
|
||||
from typing import NamedTuple
|
||||
from typing import Tuple
|
||||
from typing import Union
|
||||
|
||||
import torch
|
||||
|
||||
from funasr.metrics.common import end_detect
|
||||
from funasr.models.transformer.scorers.scorer_interface import PartialScorerInterface
|
||||
from funasr.models.transformer.scorers.scorer_interface import ScorerInterface
|
||||
|
||||
|
||||
class Hypothesis(NamedTuple):
|
||||
"""Hypothesis data type."""
|
||||
|
||||
yseq: torch.Tensor
|
||||
score: Union[float, torch.Tensor] = 0
|
||||
scores: Dict[str, Union[float, torch.Tensor]] = dict()
|
||||
states: Dict[str, Any] = dict()
|
||||
|
||||
def asdict(self) -> dict:
|
||||
"""Convert data to JSON-friendly dict."""
|
||||
return self._replace(
|
||||
yseq=self.yseq.tolist(),
|
||||
score=float(self.score),
|
||||
scores={k: float(v) for k, v in self.scores.items()},
|
||||
)._asdict()
|
||||
|
||||
|
||||
class BeamSearchScama(torch.nn.Module):
|
||||
"""Beam search implementation."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
scorers: Dict[str, ScorerInterface],
|
||||
weights: Dict[str, float],
|
||||
beam_size: int,
|
||||
vocab_size: int,
|
||||
sos: int,
|
||||
eos: int,
|
||||
token_list: List[str] = None,
|
||||
pre_beam_ratio: float = 1.5,
|
||||
pre_beam_score_key: str = None,
|
||||
):
|
||||
"""Initialize beam search.
|
||||
|
||||
Args:
|
||||
scorers (dict[str, ScorerInterface]): Dict of decoder modules
|
||||
e.g., Decoder, CTCPrefixScorer, LM
|
||||
The scorer will be ignored if it is `None`
|
||||
weights (dict[str, float]): Dict of weights for each scorers
|
||||
The scorer will be ignored if its weight is 0
|
||||
beam_size (int): The number of hypotheses kept during search
|
||||
vocab_size (int): The number of vocabulary
|
||||
sos (int): Start of sequence id
|
||||
eos (int): End of sequence id
|
||||
token_list (list[str]): List of tokens for debug log
|
||||
pre_beam_score_key (str): key of scores to perform pre-beam search
|
||||
pre_beam_ratio (float): beam size in the pre-beam search
|
||||
will be `int(pre_beam_ratio * beam_size)`
|
||||
|
||||
"""
|
||||
super().__init__()
|
||||
# set scorers
|
||||
self.weights = weights
|
||||
self.scorers = dict()
|
||||
self.full_scorers = dict()
|
||||
self.part_scorers = dict()
|
||||
# this module dict is required for recursive cast
|
||||
# `self.to(device, dtype)` in `recog.py`
|
||||
self.nn_dict = torch.nn.ModuleDict()
|
||||
for k, v in scorers.items():
|
||||
w = weights.get(k, 0)
|
||||
if w == 0 or v is None:
|
||||
continue
|
||||
assert isinstance(
|
||||
v, ScorerInterface
|
||||
), f"{k} ({type(v)}) does not implement ScorerInterface"
|
||||
self.scorers[k] = v
|
||||
if isinstance(v, PartialScorerInterface):
|
||||
self.part_scorers[k] = v
|
||||
else:
|
||||
self.full_scorers[k] = v
|
||||
if isinstance(v, torch.nn.Module):
|
||||
self.nn_dict[k] = v
|
||||
|
||||
# set configurations
|
||||
self.sos = sos
|
||||
self.eos = eos
|
||||
self.token_list = token_list
|
||||
self.pre_beam_size = int(pre_beam_ratio * beam_size)
|
||||
self.beam_size = beam_size
|
||||
self.n_vocab = vocab_size
|
||||
if (
|
||||
pre_beam_score_key is not None
|
||||
and pre_beam_score_key != "full"
|
||||
and pre_beam_score_key not in self.full_scorers
|
||||
):
|
||||
raise KeyError(f"{pre_beam_score_key} is not found in {self.full_scorers}")
|
||||
self.pre_beam_score_key = pre_beam_score_key
|
||||
self.do_pre_beam = (
|
||||
self.pre_beam_score_key is not None
|
||||
and self.pre_beam_size < self.n_vocab
|
||||
and len(self.part_scorers) > 0
|
||||
)
|
||||
|
||||
def init_hyp(self, x: torch.Tensor) -> List[Hypothesis]:
|
||||
"""Get an initial hypothesis data.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): The encoder output feature
|
||||
|
||||
Returns:
|
||||
Hypothesis: The initial hypothesis.
|
||||
|
||||
"""
|
||||
init_states = dict()
|
||||
init_scores = dict()
|
||||
for k, d in self.scorers.items():
|
||||
init_states[k] = d.init_state(x)
|
||||
init_scores[k] = 0.0
|
||||
return [
|
||||
Hypothesis(
|
||||
score=0.0,
|
||||
scores=init_scores,
|
||||
states=init_states,
|
||||
yseq=torch.tensor([self.sos], device=x.device),
|
||||
)
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def append_token(xs: torch.Tensor, x: int) -> torch.Tensor:
|
||||
"""Append new token to prefix tokens.
|
||||
|
||||
Args:
|
||||
xs (torch.Tensor): The prefix token
|
||||
x (int): The new token to append
|
||||
|
||||
Returns:
|
||||
torch.Tensor: New tensor contains: xs + [x] with xs.dtype and xs.device
|
||||
|
||||
"""
|
||||
x = torch.tensor([x], dtype=xs.dtype, device=xs.device)
|
||||
return torch.cat((xs, x))
|
||||
|
||||
def score_full(
|
||||
self,
|
||||
hyp: Hypothesis,
|
||||
x: torch.Tensor,
|
||||
x_mask: torch.Tensor = None,
|
||||
pre_acoustic_embeds: torch.Tensor = None,
|
||||
) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]:
|
||||
"""Score new hypothesis by `self.full_scorers`.
|
||||
|
||||
Args:
|
||||
hyp (Hypothesis): Hypothesis with prefix tokens to score
|
||||
x (torch.Tensor): Corresponding input feature
|
||||
|
||||
Returns:
|
||||
Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: Tuple of
|
||||
score dict of `hyp` that has string keys of `self.full_scorers`
|
||||
and tensor score values of shape: `(self.n_vocab,)`,
|
||||
and state dict that has string keys
|
||||
and state values of `self.full_scorers`
|
||||
|
||||
"""
|
||||
scores = dict()
|
||||
states = dict()
|
||||
for k, d in self.full_scorers.items():
|
||||
scores[k], states[k] = d.score(
|
||||
hyp.yseq, hyp.states[k], x, x_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds
|
||||
)
|
||||
return scores, states
|
||||
|
||||
def score_partial(
|
||||
self, hyp: Hypothesis, ids: torch.Tensor, x: torch.Tensor
|
||||
) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]:
|
||||
"""Score new hypothesis by `self.part_scorers`.
|
||||
|
||||
Args:
|
||||
hyp (Hypothesis): Hypothesis with prefix tokens to score
|
||||
ids (torch.Tensor): 1D tensor of new partial tokens to score
|
||||
x (torch.Tensor): Corresponding input feature
|
||||
|
||||
Returns:
|
||||
Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: Tuple of
|
||||
score dict of `hyp` that has string keys of `self.part_scorers`
|
||||
and tensor score values of shape: `(len(ids),)`,
|
||||
and state dict that has string keys
|
||||
and state values of `self.part_scorers`
|
||||
|
||||
"""
|
||||
scores = dict()
|
||||
states = dict()
|
||||
for k, d in self.part_scorers.items():
|
||||
scores[k], states[k] = d.score_partial(hyp.yseq, ids, hyp.states[k], x)
|
||||
return scores, states
|
||||
|
||||
def beam(
|
||||
self, weighted_scores: torch.Tensor, ids: torch.Tensor
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Compute topk full token ids and partial token ids.
|
||||
|
||||
Args:
|
||||
weighted_scores (torch.Tensor): The weighted sum scores for each tokens.
|
||||
Its shape is `(self.n_vocab,)`.
|
||||
ids (torch.Tensor): The partial token ids to compute topk
|
||||
|
||||
Returns:
|
||||
Tuple[torch.Tensor, torch.Tensor]:
|
||||
The topk full token ids and partial token ids.
|
||||
Their shapes are `(self.beam_size,)`
|
||||
|
||||
"""
|
||||
# no pre beam performed
|
||||
if weighted_scores.size(0) == ids.size(0):
|
||||
top_ids = weighted_scores.topk(self.beam_size)[1]
|
||||
return top_ids, top_ids
|
||||
|
||||
# mask pruned in pre-beam not to select in topk
|
||||
tmp = weighted_scores[ids]
|
||||
weighted_scores[:] = -float("inf")
|
||||
weighted_scores[ids] = tmp
|
||||
top_ids = weighted_scores.topk(self.beam_size)[1]
|
||||
local_ids = weighted_scores[ids].topk(self.beam_size)[1]
|
||||
return top_ids, local_ids
|
||||
|
||||
@staticmethod
|
||||
def merge_scores(
|
||||
prev_scores: Dict[str, float],
|
||||
next_full_scores: Dict[str, torch.Tensor],
|
||||
full_idx: int,
|
||||
next_part_scores: Dict[str, torch.Tensor],
|
||||
part_idx: int,
|
||||
) -> Dict[str, torch.Tensor]:
|
||||
"""Merge scores for new hypothesis.
|
||||
|
||||
Args:
|
||||
prev_scores (Dict[str, float]):
|
||||
The previous hypothesis scores by `self.scorers`
|
||||
next_full_scores (Dict[str, torch.Tensor]): scores by `self.full_scorers`
|
||||
full_idx (int): The next token id for `next_full_scores`
|
||||
next_part_scores (Dict[str, torch.Tensor]):
|
||||
scores of partial tokens by `self.part_scorers`
|
||||
part_idx (int): The new token id for `next_part_scores`
|
||||
|
||||
Returns:
|
||||
Dict[str, torch.Tensor]: The new score dict.
|
||||
Its keys are names of `self.full_scorers` and `self.part_scorers`.
|
||||
Its values are scalar tensors by the scorers.
|
||||
|
||||
"""
|
||||
new_scores = dict()
|
||||
for k, v in next_full_scores.items():
|
||||
new_scores[k] = prev_scores[k] + v[full_idx]
|
||||
for k, v in next_part_scores.items():
|
||||
new_scores[k] = prev_scores[k] + v[part_idx]
|
||||
return new_scores
|
||||
|
||||
def merge_states(self, states: Any, part_states: Any, part_idx: int) -> Any:
|
||||
"""Merge states for new hypothesis.
|
||||
|
||||
Args:
|
||||
states: states of `self.full_scorers`
|
||||
part_states: states of `self.part_scorers`
|
||||
part_idx (int): The new token id for `part_scores`
|
||||
|
||||
Returns:
|
||||
Dict[str, torch.Tensor]: The new score dict.
|
||||
Its keys are names of `self.full_scorers` and `self.part_scorers`.
|
||||
Its values are states of the scorers.
|
||||
|
||||
"""
|
||||
new_states = dict()
|
||||
for k, v in states.items():
|
||||
new_states[k] = v
|
||||
for k, d in self.part_scorers.items():
|
||||
new_states[k] = d.select_state(part_states[k], part_idx)
|
||||
return new_states
|
||||
|
||||
def search(
|
||||
self,
|
||||
running_hyps: List[Hypothesis],
|
||||
x: torch.Tensor,
|
||||
x_mask: torch.Tensor = None,
|
||||
pre_acoustic_embeds: torch.Tensor = None,
|
||||
) -> List[Hypothesis]:
|
||||
"""Search new tokens for running hypotheses and encoded speech x.
|
||||
|
||||
Args:
|
||||
running_hyps (List[Hypothesis]): Running hypotheses on beam
|
||||
x (torch.Tensor): Encoded speech feature (T, D)
|
||||
|
||||
Returns:
|
||||
List[Hypotheses]: Best sorted hypotheses
|
||||
|
||||
"""
|
||||
best_hyps = []
|
||||
part_ids = torch.arange(self.n_vocab, device=x.device) # no pre-beam
|
||||
for hyp in running_hyps:
|
||||
# scoring
|
||||
weighted_scores = torch.zeros(self.n_vocab, dtype=x.dtype, device=x.device)
|
||||
scores, states = self.score_full(
|
||||
hyp, x, x_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds
|
||||
)
|
||||
for k in self.full_scorers:
|
||||
weighted_scores += self.weights[k] * scores[k]
|
||||
# partial scoring
|
||||
if self.do_pre_beam:
|
||||
pre_beam_scores = (
|
||||
weighted_scores
|
||||
if self.pre_beam_score_key == "full"
|
||||
else scores[self.pre_beam_score_key]
|
||||
)
|
||||
part_ids = torch.topk(pre_beam_scores, self.pre_beam_size)[1]
|
||||
part_scores, part_states = self.score_partial(hyp, part_ids, x)
|
||||
for k in self.part_scorers:
|
||||
weighted_scores[part_ids] += self.weights[k] * part_scores[k]
|
||||
# add previous hyp score
|
||||
weighted_scores += hyp.score
|
||||
|
||||
# update hyps
|
||||
for j, part_j in zip(*self.beam(weighted_scores, part_ids)):
|
||||
# will be (2 x beam at most)
|
||||
best_hyps.append(
|
||||
Hypothesis(
|
||||
score=weighted_scores[j],
|
||||
yseq=self.append_token(hyp.yseq, j),
|
||||
scores=self.merge_scores(hyp.scores, scores, j, part_scores, part_j),
|
||||
states=self.merge_states(states, part_states, part_j),
|
||||
)
|
||||
)
|
||||
|
||||
# sort and prune 2 x beam -> beam
|
||||
best_hyps = sorted(best_hyps, key=lambda x: x.score, reverse=True)[
|
||||
: min(len(best_hyps), self.beam_size)
|
||||
]
|
||||
return best_hyps
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
scama_mask: torch.Tensor = None,
|
||||
pre_acoustic_embeds: torch.Tensor = None,
|
||||
maxlenratio: float = 0.0,
|
||||
minlenratio: float = 0.0,
|
||||
maxlen: int = None,
|
||||
minlen: int = 0,
|
||||
) -> List[Hypothesis]:
|
||||
"""Perform beam search.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): Encoded speech feature (T, D)
|
||||
maxlenratio (float): Input length ratio to obtain max output length.
|
||||
If maxlenratio=0.0 (default), it uses a end-detect function
|
||||
to automatically find maximum hypothesis lengths
|
||||
If maxlenratio<0.0, its absolute value is interpreted
|
||||
as a constant max output length.
|
||||
minlenratio (float): Input length ratio to obtain min output length.
|
||||
|
||||
Returns:
|
||||
list[Hypothesis]: N-best decoding results
|
||||
|
||||
"""
|
||||
if maxlen is None:
|
||||
# set length bounds
|
||||
if maxlenratio == 0:
|
||||
maxlen = x.shape[0]
|
||||
elif maxlenratio < 0:
|
||||
maxlen = -1 * int(maxlenratio)
|
||||
else:
|
||||
maxlen = max(1, int(maxlenratio * x.size(0)))
|
||||
minlen = int(minlenratio * x.size(0))
|
||||
|
||||
logging.info("decoder input length: " + str(x.shape[0]))
|
||||
logging.info("max output length: " + str(maxlen))
|
||||
logging.info("min output length: " + str(minlen))
|
||||
|
||||
# main loop of prefix search
|
||||
running_hyps = self.init_hyp(x)
|
||||
ended_hyps = []
|
||||
for i in range(maxlen):
|
||||
logging.debug("position " + str(i))
|
||||
mask_enc = None
|
||||
if scama_mask is not None:
|
||||
token_num_predictor = scama_mask.size(1)
|
||||
token_id_slice = min(i, token_num_predictor - 1)
|
||||
mask_enc = scama_mask[:, token_id_slice : token_id_slice + 1, :]
|
||||
# if mask_enc.size(1) == 0:
|
||||
# mask_enc = scama_mask[:, -2:-1, :]
|
||||
# # mask_enc = torch.zeros_like(mask_enc)
|
||||
pre_acoustic_embeds_cur = None
|
||||
if pre_acoustic_embeds is not None:
|
||||
b, t, d = pre_acoustic_embeds.size()
|
||||
pad = torch.zeros((b, 1, d), dtype=pre_acoustic_embeds.dtype).to(
|
||||
device=pre_acoustic_embeds.device
|
||||
)
|
||||
pre_acoustic_embeds = torch.cat((pre_acoustic_embeds, pad), dim=1)
|
||||
token_id_slice = min(i, t)
|
||||
pre_acoustic_embeds_cur = pre_acoustic_embeds[
|
||||
:, token_id_slice : token_id_slice + 1, :
|
||||
]
|
||||
|
||||
best = self.search(
|
||||
running_hyps, x, x_mask=mask_enc, pre_acoustic_embeds=pre_acoustic_embeds_cur
|
||||
)
|
||||
# post process of one iteration
|
||||
running_hyps = self.post_process(i, maxlen, maxlenratio, best, ended_hyps)
|
||||
# end detection
|
||||
if maxlenratio == 0.0 and end_detect([h.asdict() for h in ended_hyps], i):
|
||||
logging.info(f"end detected at {i}")
|
||||
break
|
||||
if len(running_hyps) == 0:
|
||||
logging.info("no hypothesis. Finish decoding.")
|
||||
break
|
||||
else:
|
||||
logging.debug(f"remained hypotheses: {len(running_hyps)}")
|
||||
|
||||
nbest_hyps = sorted(ended_hyps, key=lambda x: x.score, reverse=True)
|
||||
# check the number of hypotheses reaching to eos
|
||||
if len(nbest_hyps) == 0:
|
||||
logging.warning(
|
||||
"there is no N-best results, perform recognition " "again with smaller minlenratio."
|
||||
)
|
||||
return (
|
||||
[]
|
||||
if minlenratio < 0.1
|
||||
else self.forward(x, maxlenratio, max(0.0, minlenratio - 0.1))
|
||||
)
|
||||
|
||||
# report the best result
|
||||
for x in nbest_hyps:
|
||||
yseq = "".join([self.token_list[x] for x in x.yseq])
|
||||
logging.debug("nbest: y: {}, yseq: {}, score: {}".format(x.yseq, yseq, x.score))
|
||||
best = nbest_hyps[0]
|
||||
for k, v in best.scores.items():
|
||||
logging.info(f"{v:6.2f} * {self.weights[k]:3} = {v * self.weights[k]:6.2f} for {k}")
|
||||
logging.info(f"total log probability: {best.score:.2f}")
|
||||
logging.info(f"normalized log probability: {best.score / len(best.yseq):.2f}")
|
||||
logging.info(f"total number of ended hypotheses: {len(nbest_hyps)}")
|
||||
if self.token_list is not None:
|
||||
logging.info(
|
||||
"best hypo: " + "".join([self.token_list[x] for x in best.yseq[1:-1]]) + "\n"
|
||||
)
|
||||
return nbest_hyps
|
||||
|
||||
def post_process(
|
||||
self,
|
||||
i: int,
|
||||
maxlen: int,
|
||||
maxlenratio: float,
|
||||
running_hyps: List[Hypothesis],
|
||||
ended_hyps: List[Hypothesis],
|
||||
) -> List[Hypothesis]:
|
||||
"""Perform post-processing of beam search iterations.
|
||||
|
||||
Args:
|
||||
i (int): The length of hypothesis tokens.
|
||||
maxlen (int): The maximum length of tokens in beam search.
|
||||
maxlenratio (int): The maximum length ratio in beam search.
|
||||
running_hyps (List[Hypothesis]): The running hypotheses in beam search.
|
||||
ended_hyps (List[Hypothesis]): The ended hypotheses in beam search.
|
||||
|
||||
Returns:
|
||||
List[Hypothesis]: The new running hypotheses.
|
||||
|
||||
"""
|
||||
logging.debug(f"the number of running hypotheses: {len(running_hyps)}")
|
||||
if self.token_list is not None:
|
||||
logging.debug(
|
||||
"best hypo: " + "".join([self.token_list[x] for x in running_hyps[0].yseq[1:]])
|
||||
)
|
||||
# add eos in the final loop to avoid that there are no ended hyps
|
||||
if i == maxlen - 1:
|
||||
logging.info("adding <eos> in the last position in the loop")
|
||||
running_hyps = [
|
||||
h._replace(yseq=self.append_token(h.yseq, self.eos)) for h in running_hyps
|
||||
]
|
||||
|
||||
# add ended hypotheses to a final list, and removed them from current hypotheses
|
||||
# (this will be a problem, number of hyps < beam)
|
||||
remained_hyps = []
|
||||
for hyp in running_hyps:
|
||||
if hyp.yseq[-1] == self.eos:
|
||||
# e.g., Word LM needs to add final <eos> score
|
||||
for k, d in chain(self.full_scorers.items(), self.part_scorers.items()):
|
||||
s = d.final_score(hyp.states[k])
|
||||
hyp.scores[k] += s
|
||||
hyp = hyp._replace(score=hyp.score + self.weights[k] * s)
|
||||
ended_hyps.append(hyp)
|
||||
else:
|
||||
remained_hyps.append(hyp)
|
||||
return remained_hyps
|
||||
|
||||
|
||||
class BeamSearchScamaStreaming(torch.nn.Module):
|
||||
"""Beam search implementation."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
scorers: Dict[str, ScorerInterface],
|
||||
weights: Dict[str, float],
|
||||
beam_size: int,
|
||||
vocab_size: int,
|
||||
sos: int,
|
||||
eos: int,
|
||||
token_list: List[str] = None,
|
||||
pre_beam_ratio: float = 1.5,
|
||||
pre_beam_score_key: str = None,
|
||||
):
|
||||
"""Initialize beam search.
|
||||
|
||||
Args:
|
||||
scorers (dict[str, ScorerInterface]): Dict of decoder modules
|
||||
e.g., Decoder, CTCPrefixScorer, LM
|
||||
The scorer will be ignored if it is `None`
|
||||
weights (dict[str, float]): Dict of weights for each scorers
|
||||
The scorer will be ignored if its weight is 0
|
||||
beam_size (int): The number of hypotheses kept during search
|
||||
vocab_size (int): The number of vocabulary
|
||||
sos (int): Start of sequence id
|
||||
eos (int): End of sequence id
|
||||
token_list (list[str]): List of tokens for debug log
|
||||
pre_beam_score_key (str): key of scores to perform pre-beam search
|
||||
pre_beam_ratio (float): beam size in the pre-beam search
|
||||
will be `int(pre_beam_ratio * beam_size)`
|
||||
|
||||
"""
|
||||
super().__init__()
|
||||
# set scorers
|
||||
self.weights = weights
|
||||
self.scorers = dict()
|
||||
self.full_scorers = dict()
|
||||
self.part_scorers = dict()
|
||||
# this module dict is required for recursive cast
|
||||
# `self.to(device, dtype)` in `recog.py`
|
||||
self.nn_dict = torch.nn.ModuleDict()
|
||||
for k, v in scorers.items():
|
||||
w = weights.get(k, 0)
|
||||
if w == 0 or v is None:
|
||||
continue
|
||||
assert isinstance(
|
||||
v, ScorerInterface
|
||||
), f"{k} ({type(v)}) does not implement ScorerInterface"
|
||||
self.scorers[k] = v
|
||||
if isinstance(v, PartialScorerInterface):
|
||||
self.part_scorers[k] = v
|
||||
else:
|
||||
self.full_scorers[k] = v
|
||||
if isinstance(v, torch.nn.Module):
|
||||
self.nn_dict[k] = v
|
||||
|
||||
# set configurations
|
||||
self.sos = sos
|
||||
self.eos = eos
|
||||
self.token_list = token_list
|
||||
self.pre_beam_size = int(pre_beam_ratio * beam_size)
|
||||
self.beam_size = beam_size
|
||||
self.n_vocab = vocab_size
|
||||
if (
|
||||
pre_beam_score_key is not None
|
||||
and pre_beam_score_key != "full"
|
||||
and pre_beam_score_key not in self.full_scorers
|
||||
):
|
||||
raise KeyError(f"{pre_beam_score_key} is not found in {self.full_scorers}")
|
||||
self.pre_beam_score_key = pre_beam_score_key
|
||||
self.do_pre_beam = (
|
||||
self.pre_beam_score_key is not None
|
||||
and self.pre_beam_size < self.n_vocab
|
||||
and len(self.part_scorers) > 0
|
||||
)
|
||||
|
||||
def init_hyp(self, x) -> List[Hypothesis]:
|
||||
"""Get an initial hypothesis data.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): The encoder output feature
|
||||
|
||||
Returns:
|
||||
Hypothesis: The initial hypothesis.
|
||||
|
||||
"""
|
||||
init_states = dict()
|
||||
init_scores = dict()
|
||||
for k, d in self.scorers.items():
|
||||
init_states[k] = d.init_state(x)
|
||||
init_scores[k] = 0.0
|
||||
return [
|
||||
Hypothesis(
|
||||
score=0.0,
|
||||
scores=init_scores,
|
||||
states=init_states,
|
||||
yseq=torch.tensor([self.sos], device=x.device),
|
||||
)
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def append_token(xs: torch.Tensor, x: int) -> torch.Tensor:
|
||||
"""Append new token to prefix tokens.
|
||||
|
||||
Args:
|
||||
xs (torch.Tensor): The prefix token
|
||||
x (int): The new token to append
|
||||
|
||||
Returns:
|
||||
torch.Tensor: New tensor contains: xs + [x] with xs.dtype and xs.device
|
||||
|
||||
"""
|
||||
x = torch.tensor([x], dtype=xs.dtype, device=xs.device)
|
||||
return torch.cat((xs, x))
|
||||
|
||||
def score_full(
|
||||
self,
|
||||
hyp: Hypothesis,
|
||||
x: torch.Tensor,
|
||||
x_mask: torch.Tensor = None,
|
||||
pre_acoustic_embeds: torch.Tensor = None,
|
||||
cache: dict = {},
|
||||
) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]:
|
||||
"""Score new hypothesis by `self.full_scorers`.
|
||||
|
||||
Args:
|
||||
hyp (Hypothesis): Hypothesis with prefix tokens to score
|
||||
x (torch.Tensor): Corresponding input feature
|
||||
|
||||
Returns:
|
||||
Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: Tuple of
|
||||
score dict of `hyp` that has string keys of `self.full_scorers`
|
||||
and tensor score values of shape: `(self.n_vocab,)`,
|
||||
and state dict that has string keys
|
||||
and state values of `self.full_scorers`
|
||||
|
||||
"""
|
||||
scores = dict()
|
||||
states = dict()
|
||||
for k, d in self.full_scorers.items():
|
||||
scores[k], states[k] = d.score(
|
||||
hyp.yseq,
|
||||
hyp.states[k],
|
||||
x,
|
||||
x_mask=x_mask,
|
||||
pre_acoustic_embeds=pre_acoustic_embeds,
|
||||
cache=cache,
|
||||
)
|
||||
return scores, states
|
||||
|
||||
def score_partial(
|
||||
self, hyp: Hypothesis, ids: torch.Tensor, x: torch.Tensor
|
||||
) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]:
|
||||
"""Score new hypothesis by `self.part_scorers`.
|
||||
|
||||
Args:
|
||||
hyp (Hypothesis): Hypothesis with prefix tokens to score
|
||||
ids (torch.Tensor): 1D tensor of new partial tokens to score
|
||||
x (torch.Tensor): Corresponding input feature
|
||||
|
||||
Returns:
|
||||
Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: Tuple of
|
||||
score dict of `hyp` that has string keys of `self.part_scorers`
|
||||
and tensor score values of shape: `(len(ids),)`,
|
||||
and state dict that has string keys
|
||||
and state values of `self.part_scorers`
|
||||
|
||||
"""
|
||||
scores = dict()
|
||||
states = dict()
|
||||
for k, d in self.part_scorers.items():
|
||||
scores[k], states[k] = d.score_partial(hyp.yseq, ids, hyp.states[k], x)
|
||||
return scores, states
|
||||
|
||||
def beam(
|
||||
self, weighted_scores: torch.Tensor, ids: torch.Tensor
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Compute topk full token ids and partial token ids.
|
||||
|
||||
Args:
|
||||
weighted_scores (torch.Tensor): The weighted sum scores for each tokens.
|
||||
Its shape is `(self.n_vocab,)`.
|
||||
ids (torch.Tensor): The partial token ids to compute topk
|
||||
|
||||
Returns:
|
||||
Tuple[torch.Tensor, torch.Tensor]:
|
||||
The topk full token ids and partial token ids.
|
||||
Their shapes are `(self.beam_size,)`
|
||||
|
||||
"""
|
||||
# no pre beam performed
|
||||
if weighted_scores.size(0) == ids.size(0):
|
||||
top_ids = weighted_scores.topk(self.beam_size)[1]
|
||||
return top_ids, top_ids
|
||||
|
||||
# mask pruned in pre-beam not to select in topk
|
||||
tmp = weighted_scores[ids]
|
||||
weighted_scores[:] = -float("inf")
|
||||
weighted_scores[ids] = tmp
|
||||
top_ids = weighted_scores.topk(self.beam_size)[1]
|
||||
local_ids = weighted_scores[ids].topk(self.beam_size)[1]
|
||||
return top_ids, local_ids
|
||||
|
||||
@staticmethod
|
||||
def merge_scores(
|
||||
prev_scores: Dict[str, float],
|
||||
next_full_scores: Dict[str, torch.Tensor],
|
||||
full_idx: int,
|
||||
next_part_scores: Dict[str, torch.Tensor],
|
||||
part_idx: int,
|
||||
) -> Dict[str, torch.Tensor]:
|
||||
"""Merge scores for new hypothesis.
|
||||
|
||||
Args:
|
||||
prev_scores (Dict[str, float]):
|
||||
The previous hypothesis scores by `self.scorers`
|
||||
next_full_scores (Dict[str, torch.Tensor]): scores by `self.full_scorers`
|
||||
full_idx (int): The next token id for `next_full_scores`
|
||||
next_part_scores (Dict[str, torch.Tensor]):
|
||||
scores of partial tokens by `self.part_scorers`
|
||||
part_idx (int): The new token id for `next_part_scores`
|
||||
|
||||
Returns:
|
||||
Dict[str, torch.Tensor]: The new score dict.
|
||||
Its keys are names of `self.full_scorers` and `self.part_scorers`.
|
||||
Its values are scalar tensors by the scorers.
|
||||
|
||||
"""
|
||||
new_scores = dict()
|
||||
for k, v in next_full_scores.items():
|
||||
new_scores[k] = prev_scores[k] + v[full_idx]
|
||||
for k, v in next_part_scores.items():
|
||||
new_scores[k] = prev_scores[k] + v[part_idx]
|
||||
return new_scores
|
||||
|
||||
def merge_states(self, states: Any, part_states: Any, part_idx: int) -> Any:
|
||||
"""Merge states for new hypothesis.
|
||||
|
||||
Args:
|
||||
states: states of `self.full_scorers`
|
||||
part_states: states of `self.part_scorers`
|
||||
part_idx (int): The new token id for `part_scores`
|
||||
|
||||
Returns:
|
||||
Dict[str, torch.Tensor]: The new score dict.
|
||||
Its keys are names of `self.full_scorers` and `self.part_scorers`.
|
||||
Its values are states of the scorers.
|
||||
|
||||
"""
|
||||
new_states = dict()
|
||||
for k, v in states.items():
|
||||
new_states[k] = v
|
||||
for k, d in self.part_scorers.items():
|
||||
new_states[k] = d.select_state(part_states[k], part_idx)
|
||||
return new_states
|
||||
|
||||
def search(
|
||||
self,
|
||||
running_hyps: List[Hypothesis],
|
||||
x: torch.Tensor,
|
||||
x_mask: torch.Tensor = None,
|
||||
pre_acoustic_embeds: torch.Tensor = None,
|
||||
cache: dict = {},
|
||||
) -> List[Hypothesis]:
|
||||
"""Search new tokens for running hypotheses and encoded speech x.
|
||||
|
||||
Args:
|
||||
running_hyps (List[Hypothesis]): Running hypotheses on beam
|
||||
x (torch.Tensor): Encoded speech feature (T, D)
|
||||
|
||||
Returns:
|
||||
List[Hypotheses]: Best sorted hypotheses
|
||||
|
||||
"""
|
||||
best_hyps = []
|
||||
part_ids = torch.arange(self.n_vocab, device=x.device) # no pre-beam
|
||||
for hyp in running_hyps:
|
||||
# scoring
|
||||
weighted_scores = torch.zeros(self.n_vocab, dtype=x.dtype, device=x.device)
|
||||
scores, states = self.score_full(
|
||||
hyp, x, x_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds, cache=cache
|
||||
)
|
||||
for k in self.full_scorers:
|
||||
weighted_scores += self.weights[k] * scores[k]
|
||||
# partial scoring
|
||||
if self.do_pre_beam:
|
||||
pre_beam_scores = (
|
||||
weighted_scores
|
||||
if self.pre_beam_score_key == "full"
|
||||
else scores[self.pre_beam_score_key]
|
||||
)
|
||||
part_ids = torch.topk(pre_beam_scores, self.pre_beam_size)[1]
|
||||
part_scores, part_states = self.score_partial(hyp, part_ids, x)
|
||||
for k in self.part_scorers:
|
||||
weighted_scores[part_ids] += self.weights[k] * part_scores[k]
|
||||
# add previous hyp score
|
||||
weighted_scores += hyp.score
|
||||
|
||||
# update hyps
|
||||
for j, part_j in zip(*self.beam(weighted_scores, part_ids)):
|
||||
# will be (2 x beam at most)
|
||||
best_hyps.append(
|
||||
Hypothesis(
|
||||
score=weighted_scores[j],
|
||||
yseq=self.append_token(hyp.yseq, j),
|
||||
scores=self.merge_scores(hyp.scores, scores, j, part_scores, part_j),
|
||||
states=self.merge_states(states, part_states, part_j),
|
||||
)
|
||||
)
|
||||
|
||||
# sort and prune 2 x beam -> beam
|
||||
best_hyps = sorted(best_hyps, key=lambda x: x.score, reverse=True)[
|
||||
: min(len(best_hyps), self.beam_size)
|
||||
]
|
||||
return best_hyps
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
scama_mask: torch.Tensor = None,
|
||||
pre_acoustic_embeds: torch.Tensor = None,
|
||||
maxlenratio: float = 0.0,
|
||||
minlenratio: float = 0.0,
|
||||
maxlen: int = None,
|
||||
minlen: int = 0,
|
||||
cache: dict = {},
|
||||
) -> List[Hypothesis]:
|
||||
"""Perform beam search.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): Encoded speech feature (T, D)
|
||||
maxlenratio (float): Input length ratio to obtain max output length.
|
||||
If maxlenratio=0.0 (default), it uses a end-detect function
|
||||
to automatically find maximum hypothesis lengths
|
||||
If maxlenratio<0.0, its absolute value is interpreted
|
||||
as a constant max output length.
|
||||
minlenratio (float): Input length ratio to obtain min output length.
|
||||
|
||||
Returns:
|
||||
list[Hypothesis]: N-best decoding results
|
||||
|
||||
"""
|
||||
if maxlen is None:
|
||||
# set length bounds
|
||||
if maxlenratio == 0:
|
||||
maxlen = x.shape[0]
|
||||
elif maxlenratio < 0:
|
||||
maxlen = -1 * int(maxlenratio)
|
||||
else:
|
||||
maxlen = max(1, int(maxlenratio * x.size(0)))
|
||||
minlen = int(minlenratio * x.size(0))
|
||||
|
||||
logging.info("decoder input length: " + str(x.shape[0]))
|
||||
logging.info("max output length: " + str(maxlen))
|
||||
logging.info("min output length: " + str(minlen))
|
||||
|
||||
# main loop of prefix search
|
||||
# running_hyps = self.init_hyp(x)
|
||||
running_hyps = cache["running_hyps"]
|
||||
ended_hyps = []
|
||||
for i in range(maxlen):
|
||||
logging.debug("position " + str(i))
|
||||
mask_enc = None
|
||||
# if scama_mask is not None:
|
||||
# token_num_predictor = scama_mask.size(1)
|
||||
# token_id_slice = min(i, token_num_predictor-1)
|
||||
# mask_enc = scama_mask[:, token_id_slice:token_id_slice+1, :]
|
||||
# # if mask_enc.size(1) == 0:
|
||||
# # mask_enc = scama_mask[:, -2:-1, :]
|
||||
# # # mask_enc = torch.zeros_like(mask_enc)
|
||||
pre_acoustic_embeds_cur = None
|
||||
if pre_acoustic_embeds is not None:
|
||||
b, t, d = pre_acoustic_embeds.size()
|
||||
pad = torch.zeros((b, 1, d), dtype=pre_acoustic_embeds.dtype).to(
|
||||
device=pre_acoustic_embeds.device
|
||||
)
|
||||
pre_acoustic_embeds = torch.cat((pre_acoustic_embeds, pad), dim=1)
|
||||
token_id_slice = min(i, t)
|
||||
pre_acoustic_embeds_cur = pre_acoustic_embeds[
|
||||
:, token_id_slice : token_id_slice + 1, :
|
||||
]
|
||||
|
||||
best = self.search(
|
||||
running_hyps,
|
||||
x,
|
||||
x_mask=mask_enc,
|
||||
pre_acoustic_embeds=pre_acoustic_embeds_cur,
|
||||
cache=cache["decoder"],
|
||||
)
|
||||
# post process of one iteration
|
||||
running_hyps = self.post_process(i, maxlen, maxlenratio, best, ended_hyps)
|
||||
# end detection
|
||||
if maxlenratio == 0.0 and end_detect([h.asdict() for h in ended_hyps], i):
|
||||
logging.info(f"end detected at {i}")
|
||||
break
|
||||
if len(running_hyps) == 0:
|
||||
logging.info("no hypothesis. Finish decoding.")
|
||||
break
|
||||
else:
|
||||
logging.debug(f"remained hypotheses: {len(running_hyps)}")
|
||||
|
||||
nbest_hyps = sorted(ended_hyps, key=lambda x: x.score, reverse=True)
|
||||
# check the number of hypotheses reaching to eos
|
||||
if len(nbest_hyps) == 0:
|
||||
logging.warning(
|
||||
"there is no N-best results, perform recognition " "again with smaller minlenratio."
|
||||
)
|
||||
return (
|
||||
[]
|
||||
if minlenratio < 0.1
|
||||
else self.forward(x, maxlenratio, max(0.0, minlenratio - 0.1))
|
||||
)
|
||||
|
||||
# report the best result
|
||||
for x in nbest_hyps:
|
||||
yseq = "".join([self.token_list[x] for x in x.yseq])
|
||||
logging.debug("nbest: y: {}, yseq: {}, score: {}".format(x.yseq, yseq, x.score))
|
||||
best = nbest_hyps[0]
|
||||
for k, v in best.scores.items():
|
||||
logging.info(f"{v:6.2f} * {self.weights[k]:3} = {v * self.weights[k]:6.2f} for {k}")
|
||||
logging.info(f"total log probability: {best.score:.2f}")
|
||||
logging.info(f"normalized log probability: {best.score / len(best.yseq):.2f}")
|
||||
logging.info(f"total number of ended hypotheses: {len(nbest_hyps)}")
|
||||
if self.token_list is not None:
|
||||
logging.info(
|
||||
"best hypo: " + "".join([self.token_list[x] for x in best.yseq[1:-1]]) + "\n"
|
||||
)
|
||||
return nbest_hyps
|
||||
|
||||
def post_process(
|
||||
self,
|
||||
i: int,
|
||||
maxlen: int,
|
||||
maxlenratio: float,
|
||||
running_hyps: List[Hypothesis],
|
||||
ended_hyps: List[Hypothesis],
|
||||
) -> List[Hypothesis]:
|
||||
"""Perform post-processing of beam search iterations.
|
||||
|
||||
Args:
|
||||
i (int): The length of hypothesis tokens.
|
||||
maxlen (int): The maximum length of tokens in beam search.
|
||||
maxlenratio (int): The maximum length ratio in beam search.
|
||||
running_hyps (List[Hypothesis]): The running hypotheses in beam search.
|
||||
ended_hyps (List[Hypothesis]): The ended hypotheses in beam search.
|
||||
|
||||
Returns:
|
||||
List[Hypothesis]: The new running hypotheses.
|
||||
|
||||
"""
|
||||
logging.debug(f"the number of running hypotheses: {len(running_hyps)}")
|
||||
if self.token_list is not None:
|
||||
logging.debug(
|
||||
"best hypo: " + "".join([self.token_list[x] for x in running_hyps[0].yseq[1:]])
|
||||
)
|
||||
# add eos in the final loop to avoid that there are no ended hyps
|
||||
if i == maxlen - 1:
|
||||
logging.info("adding <eos> in the last position in the loop")
|
||||
running_hyps = [
|
||||
h._replace(yseq=self.append_token(h.yseq, self.eos)) for h in running_hyps
|
||||
]
|
||||
|
||||
# add ended hypotheses to a final list, and removed them from current hypotheses
|
||||
# (this will be a problem, number of hyps < beam)
|
||||
remained_hyps = []
|
||||
for hyp in running_hyps:
|
||||
if hyp.yseq[-1] == self.eos:
|
||||
# e.g., Word LM needs to add final <eos> score
|
||||
for k, d in chain(self.full_scorers.items(), self.part_scorers.items()):
|
||||
s = d.final_score(hyp.states[k])
|
||||
hyp.scores[k] += s
|
||||
hyp = hyp._replace(score=hyp.score + self.weights[k] * s)
|
||||
ended_hyps.append(hyp)
|
||||
else:
|
||||
remained_hyps.append(hyp)
|
||||
return remained_hyps
|
||||
@@ -0,0 +1,493 @@
|
||||
import math
|
||||
import torch
|
||||
import numpy as np
|
||||
import torch.nn.functional as F
|
||||
|
||||
from funasr.models.scama.utils import sequence_mask
|
||||
from funasr.models.transformer.utils.nets_utils import make_pad_mask
|
||||
|
||||
|
||||
class overlap_chunk:
|
||||
"""
|
||||
Author: Speech Lab of DAMO Academy, Alibaba Group
|
||||
San-m: Memory equipped self-attention for end-to-end speech recognition
|
||||
https://arxiv.org/abs/2006.01713
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
chunk_size: tuple = (16,),
|
||||
stride: tuple = (10,),
|
||||
pad_left: tuple = (0,),
|
||||
encoder_att_look_back_factor: tuple = (1,),
|
||||
shfit_fsmn: int = 0,
|
||||
decoder_att_look_back_factor: tuple = (1,),
|
||||
):
|
||||
|
||||
pad_left = self.check_chunk_size_args(chunk_size, pad_left)
|
||||
encoder_att_look_back_factor = self.check_chunk_size_args(
|
||||
chunk_size, encoder_att_look_back_factor
|
||||
)
|
||||
decoder_att_look_back_factor = self.check_chunk_size_args(
|
||||
chunk_size, decoder_att_look_back_factor
|
||||
)
|
||||
(
|
||||
self.chunk_size,
|
||||
self.stride,
|
||||
self.pad_left,
|
||||
self.encoder_att_look_back_factor,
|
||||
self.decoder_att_look_back_factor,
|
||||
) = (
|
||||
chunk_size,
|
||||
stride,
|
||||
pad_left,
|
||||
encoder_att_look_back_factor,
|
||||
decoder_att_look_back_factor,
|
||||
)
|
||||
self.shfit_fsmn = shfit_fsmn
|
||||
self.x_add_mask = None
|
||||
self.x_rm_mask = None
|
||||
self.x_len = None
|
||||
self.mask_shfit_chunk = None
|
||||
self.mask_chunk_predictor = None
|
||||
self.mask_att_chunk_encoder = None
|
||||
self.mask_shift_att_chunk_decoder = None
|
||||
self.chunk_outs = None
|
||||
(
|
||||
self.chunk_size_cur,
|
||||
self.stride_cur,
|
||||
self.pad_left_cur,
|
||||
self.encoder_att_look_back_factor_cur,
|
||||
self.chunk_size_pad_shift_cur,
|
||||
) = (None, None, None, None, None)
|
||||
|
||||
def check_chunk_size_args(self, chunk_size, x):
|
||||
if len(x) < len(chunk_size):
|
||||
x = [x[0] for i in chunk_size]
|
||||
return x
|
||||
|
||||
def get_chunk_size(self, ind: int = 0):
|
||||
# with torch.no_grad:
|
||||
chunk_size, stride, pad_left, encoder_att_look_back_factor, decoder_att_look_back_factor = (
|
||||
self.chunk_size[ind],
|
||||
self.stride[ind],
|
||||
self.pad_left[ind],
|
||||
self.encoder_att_look_back_factor[ind],
|
||||
self.decoder_att_look_back_factor[ind],
|
||||
)
|
||||
(
|
||||
self.chunk_size_cur,
|
||||
self.stride_cur,
|
||||
self.pad_left_cur,
|
||||
self.encoder_att_look_back_factor_cur,
|
||||
self.chunk_size_pad_shift_cur,
|
||||
self.decoder_att_look_back_factor_cur,
|
||||
) = (
|
||||
chunk_size,
|
||||
stride,
|
||||
pad_left,
|
||||
encoder_att_look_back_factor,
|
||||
chunk_size + self.shfit_fsmn,
|
||||
decoder_att_look_back_factor,
|
||||
)
|
||||
return (
|
||||
self.chunk_size_cur,
|
||||
self.stride_cur,
|
||||
self.pad_left_cur,
|
||||
self.encoder_att_look_back_factor_cur,
|
||||
self.chunk_size_pad_shift_cur,
|
||||
)
|
||||
|
||||
def random_choice(self, training=True, decoding_ind=None):
|
||||
chunk_num = len(self.chunk_size)
|
||||
ind = 0
|
||||
if training and chunk_num > 1:
|
||||
ind = torch.randint(0, chunk_num, ()).cpu().item()
|
||||
if not training and decoding_ind is not None:
|
||||
ind = int(decoding_ind)
|
||||
|
||||
return ind
|
||||
|
||||
def gen_chunk_mask(self, x_len, ind=0, num_units=1, num_units_predictor=1):
|
||||
|
||||
with torch.no_grad():
|
||||
x_len = x_len.cpu().numpy()
|
||||
x_len_max = x_len.max()
|
||||
|
||||
chunk_size, stride, pad_left, encoder_att_look_back_factor, chunk_size_pad_shift = (
|
||||
self.get_chunk_size(ind)
|
||||
)
|
||||
shfit_fsmn = self.shfit_fsmn
|
||||
pad_right = chunk_size - stride - pad_left
|
||||
|
||||
chunk_num_batch = np.ceil(x_len / stride).astype(np.int32)
|
||||
x_len_chunk = (
|
||||
(chunk_num_batch - 1) * chunk_size_pad_shift
|
||||
+ shfit_fsmn
|
||||
+ pad_left
|
||||
+ 0
|
||||
+ x_len
|
||||
- (chunk_num_batch - 1) * stride
|
||||
)
|
||||
x_len_chunk = x_len_chunk.astype(x_len.dtype)
|
||||
x_len_chunk_max = x_len_chunk.max()
|
||||
|
||||
chunk_num = int(math.ceil(x_len_max / stride))
|
||||
dtype = np.int32
|
||||
max_len_for_x_mask_tmp = max(chunk_size, x_len_max + pad_left)
|
||||
x_add_mask = np.zeros([0, max_len_for_x_mask_tmp], dtype=dtype)
|
||||
x_rm_mask = np.zeros([max_len_for_x_mask_tmp, 0], dtype=dtype)
|
||||
mask_shfit_chunk = np.zeros([0, num_units], dtype=dtype)
|
||||
mask_chunk_predictor = np.zeros([0, num_units_predictor], dtype=dtype)
|
||||
mask_shift_att_chunk_decoder = np.zeros([0, 1], dtype=dtype)
|
||||
mask_att_chunk_encoder = np.zeros([0, chunk_num * chunk_size_pad_shift], dtype=dtype)
|
||||
for chunk_ids in range(chunk_num):
|
||||
# x_mask add
|
||||
fsmn_padding = np.zeros((shfit_fsmn, max_len_for_x_mask_tmp), dtype=dtype)
|
||||
x_mask_cur = np.diag(np.ones(chunk_size, dtype=np.float32))
|
||||
x_mask_pad_left = np.zeros((chunk_size, chunk_ids * stride), dtype=dtype)
|
||||
x_mask_pad_right = np.zeros((chunk_size, max_len_for_x_mask_tmp), dtype=dtype)
|
||||
x_cur_pad = np.concatenate([x_mask_pad_left, x_mask_cur, x_mask_pad_right], axis=1)
|
||||
x_cur_pad = x_cur_pad[:chunk_size, :max_len_for_x_mask_tmp]
|
||||
x_add_mask_fsmn = np.concatenate([fsmn_padding, x_cur_pad], axis=0)
|
||||
x_add_mask = np.concatenate([x_add_mask, x_add_mask_fsmn], axis=0)
|
||||
|
||||
# x_mask rm
|
||||
fsmn_padding = np.zeros((max_len_for_x_mask_tmp, shfit_fsmn), dtype=dtype)
|
||||
padding_mask_left = np.zeros((max_len_for_x_mask_tmp, pad_left), dtype=dtype)
|
||||
padding_mask_right = np.zeros((max_len_for_x_mask_tmp, pad_right), dtype=dtype)
|
||||
x_mask_cur = np.diag(np.ones(stride, dtype=dtype))
|
||||
x_mask_cur_pad_top = np.zeros((chunk_ids * stride, stride), dtype=dtype)
|
||||
x_mask_cur_pad_bottom = np.zeros((max_len_for_x_mask_tmp, stride), dtype=dtype)
|
||||
x_rm_mask_cur = np.concatenate(
|
||||
[x_mask_cur_pad_top, x_mask_cur, x_mask_cur_pad_bottom], axis=0
|
||||
)
|
||||
x_rm_mask_cur = x_rm_mask_cur[:max_len_for_x_mask_tmp, :stride]
|
||||
x_rm_mask_cur_fsmn = np.concatenate(
|
||||
[fsmn_padding, padding_mask_left, x_rm_mask_cur, padding_mask_right], axis=1
|
||||
)
|
||||
x_rm_mask = np.concatenate([x_rm_mask, x_rm_mask_cur_fsmn], axis=1)
|
||||
|
||||
# fsmn_padding_mask
|
||||
pad_shfit_mask = np.zeros([shfit_fsmn, num_units], dtype=dtype)
|
||||
ones_1 = np.ones([chunk_size, num_units], dtype=dtype)
|
||||
mask_shfit_chunk_cur = np.concatenate([pad_shfit_mask, ones_1], axis=0)
|
||||
mask_shfit_chunk = np.concatenate([mask_shfit_chunk, mask_shfit_chunk_cur], axis=0)
|
||||
|
||||
# predictor mask
|
||||
zeros_1 = np.zeros([shfit_fsmn + pad_left, num_units_predictor], dtype=dtype)
|
||||
ones_2 = np.ones([stride, num_units_predictor], dtype=dtype)
|
||||
zeros_3 = np.zeros(
|
||||
[chunk_size - stride - pad_left, num_units_predictor], dtype=dtype
|
||||
)
|
||||
ones_zeros = np.concatenate([ones_2, zeros_3], axis=0)
|
||||
mask_chunk_predictor_cur = np.concatenate([zeros_1, ones_zeros], axis=0)
|
||||
mask_chunk_predictor = np.concatenate(
|
||||
[mask_chunk_predictor, mask_chunk_predictor_cur], axis=0
|
||||
)
|
||||
|
||||
# encoder att mask
|
||||
zeros_1_top = np.zeros([shfit_fsmn, chunk_num * chunk_size_pad_shift], dtype=dtype)
|
||||
|
||||
zeros_2_num = max(chunk_ids - encoder_att_look_back_factor, 0)
|
||||
zeros_2 = np.zeros([chunk_size, zeros_2_num * chunk_size_pad_shift], dtype=dtype)
|
||||
|
||||
encoder_att_look_back_num = max(chunk_ids - zeros_2_num, 0)
|
||||
zeros_2_left = np.zeros([chunk_size, shfit_fsmn], dtype=dtype)
|
||||
ones_2_mid = np.ones([stride, stride], dtype=dtype)
|
||||
zeros_2_bottom = np.zeros([chunk_size - stride, stride], dtype=dtype)
|
||||
zeros_2_right = np.zeros([chunk_size, chunk_size - stride], dtype=dtype)
|
||||
ones_2 = np.concatenate([ones_2_mid, zeros_2_bottom], axis=0)
|
||||
ones_2 = np.concatenate([zeros_2_left, ones_2, zeros_2_right], axis=1)
|
||||
ones_2 = np.tile(ones_2, [1, encoder_att_look_back_num])
|
||||
|
||||
zeros_3_left = np.zeros([chunk_size, shfit_fsmn], dtype=dtype)
|
||||
ones_3_right = np.ones([chunk_size, chunk_size], dtype=dtype)
|
||||
ones_3 = np.concatenate([zeros_3_left, ones_3_right], axis=1)
|
||||
|
||||
zeros_remain_num = max(chunk_num - 1 - chunk_ids, 0)
|
||||
zeros_remain = np.zeros(
|
||||
[chunk_size, zeros_remain_num * chunk_size_pad_shift], dtype=dtype
|
||||
)
|
||||
|
||||
ones2_bottom = np.concatenate([zeros_2, ones_2, ones_3, zeros_remain], axis=1)
|
||||
mask_att_chunk_encoder_cur = np.concatenate([zeros_1_top, ones2_bottom], axis=0)
|
||||
mask_att_chunk_encoder = np.concatenate(
|
||||
[mask_att_chunk_encoder, mask_att_chunk_encoder_cur], axis=0
|
||||
)
|
||||
|
||||
# decoder fsmn_shift_att_mask
|
||||
zeros_1 = np.zeros([shfit_fsmn, 1])
|
||||
ones_1 = np.ones([chunk_size, 1])
|
||||
mask_shift_att_chunk_decoder_cur = np.concatenate([zeros_1, ones_1], axis=0)
|
||||
mask_shift_att_chunk_decoder = np.concatenate(
|
||||
[mask_shift_att_chunk_decoder, mask_shift_att_chunk_decoder_cur], axis=0
|
||||
)
|
||||
|
||||
self.x_add_mask = x_add_mask[:x_len_chunk_max, : x_len_max + pad_left]
|
||||
self.x_len_chunk = x_len_chunk
|
||||
self.x_rm_mask = x_rm_mask[:x_len_max, :x_len_chunk_max]
|
||||
self.x_len = x_len
|
||||
self.mask_shfit_chunk = mask_shfit_chunk[:x_len_chunk_max, :]
|
||||
self.mask_chunk_predictor = mask_chunk_predictor[:x_len_chunk_max, :]
|
||||
self.mask_att_chunk_encoder = mask_att_chunk_encoder[:x_len_chunk_max, :x_len_chunk_max]
|
||||
self.mask_shift_att_chunk_decoder = mask_shift_att_chunk_decoder[:x_len_chunk_max, :]
|
||||
self.chunk_outs = (
|
||||
self.x_add_mask,
|
||||
self.x_len_chunk,
|
||||
self.x_rm_mask,
|
||||
self.x_len,
|
||||
self.mask_shfit_chunk,
|
||||
self.mask_chunk_predictor,
|
||||
self.mask_att_chunk_encoder,
|
||||
self.mask_shift_att_chunk_decoder,
|
||||
)
|
||||
|
||||
return self.chunk_outs
|
||||
|
||||
def split_chunk(self, x, x_len, chunk_outs):
|
||||
"""
|
||||
:param x: (b, t, d)
|
||||
:param x_length: (b)
|
||||
:param ind: int
|
||||
:return:
|
||||
"""
|
||||
x = x[:, : x_len.max(), :]
|
||||
b, t, d = x.size()
|
||||
x_len_mask = (~make_pad_mask(x_len, maxlen=t)).to(x.device)
|
||||
x *= x_len_mask[:, :, None]
|
||||
|
||||
x_add_mask = self.get_x_add_mask(chunk_outs, x.device, dtype=x.dtype)
|
||||
x_len_chunk = self.get_x_len_chunk(chunk_outs, x_len.device, dtype=x_len.dtype)
|
||||
pad = (0, 0, self.pad_left_cur, 0)
|
||||
x = F.pad(x, pad, "constant", 0.0)
|
||||
b, t, d = x.size()
|
||||
x = torch.transpose(x, 1, 0)
|
||||
x = torch.reshape(x, [t, -1])
|
||||
x_chunk = torch.mm(x_add_mask, x)
|
||||
x_chunk = torch.reshape(x_chunk, [-1, b, d]).transpose(1, 0)
|
||||
|
||||
return x_chunk, x_len_chunk
|
||||
|
||||
def remove_chunk(self, x_chunk, x_len_chunk, chunk_outs):
|
||||
x_chunk = x_chunk[:, : x_len_chunk.max(), :]
|
||||
b, t, d = x_chunk.size()
|
||||
x_len_chunk_mask = (~make_pad_mask(x_len_chunk, maxlen=t)).to(x_chunk.device)
|
||||
x_chunk *= x_len_chunk_mask[:, :, None]
|
||||
|
||||
x_rm_mask = self.get_x_rm_mask(chunk_outs, x_chunk.device, dtype=x_chunk.dtype)
|
||||
x_len = self.get_x_len(chunk_outs, x_len_chunk.device, dtype=x_len_chunk.dtype)
|
||||
x_chunk = torch.transpose(x_chunk, 1, 0)
|
||||
x_chunk = torch.reshape(x_chunk, [t, -1])
|
||||
x = torch.mm(x_rm_mask, x_chunk)
|
||||
x = torch.reshape(x, [-1, b, d]).transpose(1, 0)
|
||||
|
||||
return x, x_len
|
||||
|
||||
def get_x_add_mask(self, chunk_outs=None, device="cpu", idx=0, dtype=torch.float32):
|
||||
with torch.no_grad():
|
||||
x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx]
|
||||
x = torch.from_numpy(x).type(dtype).to(device)
|
||||
return x
|
||||
|
||||
def get_x_len_chunk(self, chunk_outs=None, device="cpu", idx=1, dtype=torch.float32):
|
||||
with torch.no_grad():
|
||||
x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx]
|
||||
x = torch.from_numpy(x).type(dtype).to(device)
|
||||
return x
|
||||
|
||||
def get_x_rm_mask(self, chunk_outs=None, device="cpu", idx=2, dtype=torch.float32):
|
||||
with torch.no_grad():
|
||||
x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx]
|
||||
x = torch.from_numpy(x).type(dtype).to(device)
|
||||
return x
|
||||
|
||||
def get_x_len(self, chunk_outs=None, device="cpu", idx=3, dtype=torch.float32):
|
||||
with torch.no_grad():
|
||||
x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx]
|
||||
x = torch.from_numpy(x).type(dtype).to(device)
|
||||
return x
|
||||
|
||||
def get_mask_shfit_chunk(
|
||||
self, chunk_outs=None, device="cpu", batch_size=1, num_units=1, idx=4, dtype=torch.float32
|
||||
):
|
||||
with torch.no_grad():
|
||||
x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx]
|
||||
x = np.tile(
|
||||
x[
|
||||
None,
|
||||
:,
|
||||
:,
|
||||
],
|
||||
[batch_size, 1, num_units],
|
||||
)
|
||||
x = torch.from_numpy(x).type(dtype).to(device)
|
||||
return x
|
||||
|
||||
def get_mask_chunk_predictor(
|
||||
self, chunk_outs=None, device="cpu", batch_size=1, num_units=1, idx=5, dtype=torch.float32
|
||||
):
|
||||
with torch.no_grad():
|
||||
x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx]
|
||||
x = np.tile(
|
||||
x[
|
||||
None,
|
||||
:,
|
||||
:,
|
||||
],
|
||||
[batch_size, 1, num_units],
|
||||
)
|
||||
x = torch.from_numpy(x).type(dtype).to(device)
|
||||
return x
|
||||
|
||||
def get_mask_att_chunk_encoder(
|
||||
self, chunk_outs=None, device="cpu", batch_size=1, idx=6, dtype=torch.float32
|
||||
):
|
||||
with torch.no_grad():
|
||||
x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx]
|
||||
x = np.tile(
|
||||
x[
|
||||
None,
|
||||
:,
|
||||
:,
|
||||
],
|
||||
[batch_size, 1, 1],
|
||||
)
|
||||
x = torch.from_numpy(x).type(dtype).to(device)
|
||||
return x
|
||||
|
||||
def get_mask_shift_att_chunk_decoder(
|
||||
self, chunk_outs=None, device="cpu", batch_size=1, idx=7, dtype=torch.float32
|
||||
):
|
||||
with torch.no_grad():
|
||||
x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx]
|
||||
x = np.tile(x[None, None, :, 0], [batch_size, 1, 1])
|
||||
x = torch.from_numpy(x).type(dtype).to(device)
|
||||
return x
|
||||
|
||||
|
||||
def build_scama_mask_for_cross_attention_decoder(
|
||||
predictor_alignments: torch.Tensor,
|
||||
encoder_sequence_length: torch.Tensor,
|
||||
chunk_size: int = 5,
|
||||
encoder_chunk_size: int = 5,
|
||||
attention_chunk_center_bias: int = 0,
|
||||
attention_chunk_size: int = 1,
|
||||
attention_chunk_type: str = "chunk",
|
||||
step=None,
|
||||
predictor_mask_chunk_hopping: torch.Tensor = None,
|
||||
decoder_att_look_back_factor: int = 1,
|
||||
mask_shift_att_chunk_decoder: torch.Tensor = None,
|
||||
target_length: torch.Tensor = None,
|
||||
is_training=True,
|
||||
dtype: torch.dtype = torch.float32,
|
||||
):
|
||||
with torch.no_grad():
|
||||
device = predictor_alignments.device
|
||||
batch_size, chunk_num = predictor_alignments.size()
|
||||
maximum_encoder_length = encoder_sequence_length.max().item()
|
||||
int_type = predictor_alignments.dtype
|
||||
if not is_training:
|
||||
target_length = predictor_alignments.sum(dim=-1).type(encoder_sequence_length.dtype)
|
||||
maximum_target_length = target_length.max()
|
||||
predictor_alignments_cumsum = torch.cumsum(predictor_alignments, dim=1)
|
||||
predictor_alignments_cumsum = predictor_alignments_cumsum[:, None, :].repeat(
|
||||
1, maximum_target_length, 1
|
||||
)
|
||||
|
||||
index = torch.ones([batch_size, maximum_target_length], dtype=int_type).to(device)
|
||||
index = torch.cumsum(index, dim=1)
|
||||
index = index[:, :, None].repeat(1, 1, chunk_num)
|
||||
|
||||
index_div = torch.floor(torch.divide(predictor_alignments_cumsum, index)).type(int_type)
|
||||
index_div_bool_zeros = index_div == 0
|
||||
index_div_bool_zeros_count = torch.sum(index_div_bool_zeros.type(int_type), dim=-1) + 1
|
||||
|
||||
index_div_bool_zeros_count = torch.clip(index_div_bool_zeros_count, min=1, max=chunk_num)
|
||||
|
||||
index_div_bool_zeros_count *= chunk_size
|
||||
index_div_bool_zeros_count += attention_chunk_center_bias
|
||||
index_div_bool_zeros_count = torch.clip(
|
||||
index_div_bool_zeros_count - 1, min=0, max=maximum_encoder_length
|
||||
)
|
||||
index_div_bool_zeros_count_ori = index_div_bool_zeros_count
|
||||
|
||||
index_div_bool_zeros_count = (
|
||||
torch.floor(index_div_bool_zeros_count / encoder_chunk_size) + 1
|
||||
) * encoder_chunk_size
|
||||
max_len_chunk = math.ceil(maximum_encoder_length / encoder_chunk_size) * encoder_chunk_size
|
||||
|
||||
mask_flip, mask_flip2 = None, None
|
||||
if attention_chunk_size is not None:
|
||||
index_div_bool_zeros_count_beg = index_div_bool_zeros_count - attention_chunk_size
|
||||
index_div_bool_zeros_count_beg = torch.clip(
|
||||
index_div_bool_zeros_count_beg, 0, max_len_chunk
|
||||
)
|
||||
index_div_bool_zeros_count_beg_mask = sequence_mask(
|
||||
index_div_bool_zeros_count_beg, maxlen=max_len_chunk, dtype=int_type, device=device
|
||||
)
|
||||
mask_flip = 1 - index_div_bool_zeros_count_beg_mask
|
||||
attention_chunk_size2 = attention_chunk_size * (decoder_att_look_back_factor + 1)
|
||||
index_div_bool_zeros_count_beg = index_div_bool_zeros_count - attention_chunk_size2
|
||||
|
||||
index_div_bool_zeros_count_beg = torch.clip(
|
||||
index_div_bool_zeros_count_beg, 0, max_len_chunk
|
||||
)
|
||||
index_div_bool_zeros_count_beg_mask = sequence_mask(
|
||||
index_div_bool_zeros_count_beg, maxlen=max_len_chunk, dtype=int_type, device=device
|
||||
)
|
||||
mask_flip2 = 1 - index_div_bool_zeros_count_beg_mask
|
||||
|
||||
mask = sequence_mask(
|
||||
index_div_bool_zeros_count, maxlen=max_len_chunk, dtype=dtype, device=device
|
||||
)
|
||||
|
||||
if predictor_mask_chunk_hopping is not None:
|
||||
b, k, t = mask.size()
|
||||
predictor_mask_chunk_hopping = predictor_mask_chunk_hopping[:, None, :, 0].repeat(
|
||||
1, k, 1
|
||||
)
|
||||
|
||||
mask_mask_flip = mask
|
||||
if mask_flip is not None:
|
||||
mask_mask_flip = mask_flip * mask
|
||||
|
||||
def _fn():
|
||||
mask_sliced = mask[:b, :k, encoder_chunk_size:t]
|
||||
zero_pad_right = torch.zeros(
|
||||
[b, k, encoder_chunk_size], dtype=mask_sliced.dtype
|
||||
).to(device)
|
||||
mask_sliced = torch.cat([mask_sliced, zero_pad_right], dim=2)
|
||||
_, _, tt = predictor_mask_chunk_hopping.size()
|
||||
pad_right_p = max_len_chunk - tt
|
||||
predictor_mask_chunk_hopping_pad = torch.nn.functional.pad(
|
||||
predictor_mask_chunk_hopping, [0, pad_right_p], "constant", 0
|
||||
)
|
||||
masked = mask_sliced * predictor_mask_chunk_hopping_pad
|
||||
|
||||
mask_true = mask_mask_flip + masked
|
||||
return mask_true
|
||||
|
||||
mask = _fn() if t > chunk_size else mask_mask_flip
|
||||
|
||||
if mask_flip2 is not None:
|
||||
mask *= mask_flip2
|
||||
|
||||
mask_target = sequence_mask(
|
||||
target_length, maxlen=maximum_target_length, dtype=mask.dtype, device=device
|
||||
)
|
||||
mask = mask[:, :maximum_target_length, :] * mask_target[:, :, None]
|
||||
|
||||
mask_len = sequence_mask(
|
||||
encoder_sequence_length, maxlen=maximum_encoder_length, dtype=mask.dtype, device=device
|
||||
)
|
||||
mask = mask[:, :, :maximum_encoder_length] * mask_len[:, None, :]
|
||||
|
||||
if attention_chunk_type == "full":
|
||||
mask = torch.ones_like(mask).to(device)
|
||||
if mask_shift_att_chunk_decoder is not None:
|
||||
mask = mask * mask_shift_att_chunk_decoder
|
||||
mask = mask[:, :maximum_target_length, :maximum_encoder_length].type(dtype).to(device)
|
||||
|
||||
return mask
|
||||
@@ -0,0 +1,498 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- encoding: utf-8 -*-
|
||||
# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
|
||||
# MIT License (https://opensource.org/licenses/MIT)
|
||||
|
||||
from typing import List
|
||||
from typing import Tuple
|
||||
import logging
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import numpy as np
|
||||
|
||||
from funasr.models.scama import utils as myutils
|
||||
from funasr.models.transformer.decoder import BaseTransformerDecoder
|
||||
|
||||
from funasr.models.sanm.attention import (
|
||||
MultiHeadedAttentionSANMDecoder,
|
||||
MultiHeadedAttentionCrossAtt,
|
||||
)
|
||||
from funasr.models.transformer.embedding import PositionalEncoding
|
||||
from funasr.models.transformer.layer_norm import LayerNorm
|
||||
from funasr.models.sanm.positionwise_feed_forward import PositionwiseFeedForwardDecoderSANM
|
||||
from funasr.models.transformer.utils.repeat import repeat
|
||||
|
||||
from funasr.register import tables
|
||||
|
||||
|
||||
class DecoderLayerSANM(nn.Module):
|
||||
"""Single decoder layer module.
|
||||
|
||||
Args:
|
||||
size (int): Input dimension.
|
||||
self_attn (torch.nn.Module): Self-attention module instance.
|
||||
`MultiHeadedAttention` instance can be used as the argument.
|
||||
src_attn (torch.nn.Module): Self-attention module instance.
|
||||
`MultiHeadedAttention` instance can be used as the argument.
|
||||
feed_forward (torch.nn.Module): Feed-forward module instance.
|
||||
`PositionwiseFeedForward`, `MultiLayeredConv1d`, or `Conv1dLinear` instance
|
||||
can be used as the argument.
|
||||
dropout_rate (float): Dropout rate.
|
||||
normalize_before (bool): Whether to use layer_norm before the first block.
|
||||
concat_after (bool): Whether to concat attention layer's input and output.
|
||||
if True, additional linear will be applied.
|
||||
i.e. x -> x + linear(concat(x, att(x)))
|
||||
if False, no additional linear will be applied. i.e. x -> x + att(x)
|
||||
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
size,
|
||||
self_attn,
|
||||
src_attn,
|
||||
feed_forward,
|
||||
dropout_rate,
|
||||
normalize_before=True,
|
||||
concat_after=False,
|
||||
):
|
||||
"""Construct an DecoderLayer object."""
|
||||
super(DecoderLayerSANM, self).__init__()
|
||||
self.size = size
|
||||
self.self_attn = self_attn
|
||||
self.src_attn = src_attn
|
||||
self.feed_forward = feed_forward
|
||||
self.norm1 = LayerNorm(size)
|
||||
if self_attn is not None:
|
||||
self.norm2 = LayerNorm(size)
|
||||
if src_attn is not None:
|
||||
self.norm3 = LayerNorm(size)
|
||||
self.dropout = nn.Dropout(dropout_rate)
|
||||
self.normalize_before = normalize_before
|
||||
self.concat_after = concat_after
|
||||
if self.concat_after:
|
||||
self.concat_linear1 = nn.Linear(size + size, size)
|
||||
self.concat_linear2 = nn.Linear(size + size, size)
|
||||
|
||||
def forward(self, tgt, tgt_mask, memory, memory_mask=None, cache=None):
|
||||
"""Compute decoded features.
|
||||
|
||||
Args:
|
||||
tgt (torch.Tensor): Input tensor (#batch, maxlen_out, size).
|
||||
tgt_mask (torch.Tensor): Mask for input tensor (#batch, maxlen_out).
|
||||
memory (torch.Tensor): Encoded memory, float32 (#batch, maxlen_in, size).
|
||||
memory_mask (torch.Tensor): Encoded memory mask (#batch, maxlen_in).
|
||||
cache (List[torch.Tensor]): List of cached tensors.
|
||||
Each tensor shape should be (#batch, maxlen_out - 1, size).
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Output tensor(#batch, maxlen_out, size).
|
||||
torch.Tensor: Mask for output tensor (#batch, maxlen_out).
|
||||
torch.Tensor: Encoded memory (#batch, maxlen_in, size).
|
||||
torch.Tensor: Encoded memory mask (#batch, maxlen_in).
|
||||
|
||||
"""
|
||||
# tgt = self.dropout(tgt)
|
||||
residual = tgt
|
||||
if self.normalize_before:
|
||||
tgt = self.norm1(tgt)
|
||||
tgt = self.feed_forward(tgt)
|
||||
|
||||
x = tgt
|
||||
if self.self_attn:
|
||||
if self.normalize_before:
|
||||
tgt = self.norm2(tgt)
|
||||
x, _ = self.self_attn(tgt, tgt_mask)
|
||||
x = residual + self.dropout(x)
|
||||
|
||||
if self.src_attn is not None:
|
||||
residual = x
|
||||
if self.normalize_before:
|
||||
x = self.norm3(x)
|
||||
|
||||
x = residual + self.dropout(self.src_attn(x, memory, memory_mask))
|
||||
|
||||
return x, tgt_mask, memory, memory_mask, cache
|
||||
|
||||
def forward_one_step(self, tgt, tgt_mask, memory, memory_mask=None, cache=None):
|
||||
"""Compute decoded features.
|
||||
|
||||
Args:
|
||||
tgt (torch.Tensor): Input tensor (#batch, maxlen_out, size).
|
||||
tgt_mask (torch.Tensor): Mask for input tensor (#batch, maxlen_out).
|
||||
memory (torch.Tensor): Encoded memory, float32 (#batch, maxlen_in, size).
|
||||
memory_mask (torch.Tensor): Encoded memory mask (#batch, maxlen_in).
|
||||
cache (List[torch.Tensor]): List of cached tensors.
|
||||
Each tensor shape should be (#batch, maxlen_out - 1, size).
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Output tensor(#batch, maxlen_out, size).
|
||||
torch.Tensor: Mask for output tensor (#batch, maxlen_out).
|
||||
torch.Tensor: Encoded memory (#batch, maxlen_in, size).
|
||||
torch.Tensor: Encoded memory mask (#batch, maxlen_in).
|
||||
|
||||
"""
|
||||
# tgt = self.dropout(tgt)
|
||||
residual = tgt
|
||||
if self.normalize_before:
|
||||
tgt = self.norm1(tgt)
|
||||
tgt = self.feed_forward(tgt)
|
||||
|
||||
x = tgt
|
||||
if self.self_attn:
|
||||
if self.normalize_before:
|
||||
tgt = self.norm2(tgt)
|
||||
if self.training:
|
||||
cache = None
|
||||
x, cache = self.self_attn(tgt, tgt_mask, cache=cache)
|
||||
x = residual + self.dropout(x)
|
||||
|
||||
if self.src_attn is not None:
|
||||
residual = x
|
||||
if self.normalize_before:
|
||||
x = self.norm3(x)
|
||||
|
||||
x = residual + self.dropout(self.src_attn(x, memory, memory_mask))
|
||||
|
||||
return x, tgt_mask, memory, memory_mask, cache
|
||||
|
||||
def forward_chunk(
|
||||
self, tgt, memory, fsmn_cache=None, opt_cache=None, chunk_size=None, look_back=0
|
||||
):
|
||||
"""Compute decoded features.
|
||||
|
||||
Args:
|
||||
tgt (torch.Tensor): Input tensor (#batch, maxlen_out, size).
|
||||
tgt_mask (torch.Tensor): Mask for input tensor (#batch, maxlen_out).
|
||||
memory (torch.Tensor): Encoded memory, float32 (#batch, maxlen_in, size).
|
||||
memory_mask (torch.Tensor): Encoded memory mask (#batch, maxlen_in).
|
||||
cache (List[torch.Tensor]): List of cached tensors.
|
||||
Each tensor shape should be (#batch, maxlen_out - 1, size).
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Output tensor(#batch, maxlen_out, size).
|
||||
torch.Tensor: Mask for output tensor (#batch, maxlen_out).
|
||||
torch.Tensor: Encoded memory (#batch, maxlen_in, size).
|
||||
torch.Tensor: Encoded memory mask (#batch, maxlen_in).
|
||||
|
||||
"""
|
||||
residual = tgt
|
||||
if self.normalize_before:
|
||||
tgt = self.norm1(tgt)
|
||||
tgt = self.feed_forward(tgt)
|
||||
|
||||
x = tgt
|
||||
if self.self_attn:
|
||||
if self.normalize_before:
|
||||
tgt = self.norm2(tgt)
|
||||
x, fsmn_cache = self.self_attn(tgt, None, fsmn_cache)
|
||||
x = residual + self.dropout(x)
|
||||
|
||||
if self.src_attn is not None:
|
||||
residual = x
|
||||
if self.normalize_before:
|
||||
x = self.norm3(x)
|
||||
|
||||
x, opt_cache = self.src_attn.forward_chunk(x, memory, opt_cache, chunk_size, look_back)
|
||||
x = residual + x
|
||||
|
||||
return x, memory, fsmn_cache, opt_cache
|
||||
|
||||
|
||||
@tables.register("decoder_classes", "FsmnDecoderSCAMAOpt")
|
||||
class FsmnDecoderSCAMAOpt(BaseTransformerDecoder):
|
||||
"""
|
||||
Author: Shiliang Zhang, Zhifu Gao, Haoneng Luo, Ming Lei, Jie Gao, Zhijie Yan, Lei Xie
|
||||
SCAMA: Streaming chunk-aware multihead attention for online end-to-end speech recognition
|
||||
https://arxiv.org/abs/2006.01712
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_size: int,
|
||||
encoder_output_size: int,
|
||||
attention_heads: int = 4,
|
||||
linear_units: int = 2048,
|
||||
num_blocks: int = 6,
|
||||
dropout_rate: float = 0.1,
|
||||
positional_dropout_rate: float = 0.1,
|
||||
self_attention_dropout_rate: float = 0.0,
|
||||
src_attention_dropout_rate: float = 0.0,
|
||||
input_layer: str = "embed",
|
||||
use_output_layer: bool = True,
|
||||
pos_enc_class=PositionalEncoding,
|
||||
normalize_before: bool = True,
|
||||
concat_after: bool = False,
|
||||
att_layer_num: int = 6,
|
||||
kernel_size: int = 21,
|
||||
sanm_shfit: int = None,
|
||||
concat_embeds: bool = False,
|
||||
attention_dim: int = None,
|
||||
tf2torch_tensor_name_prefix_torch: str = "decoder",
|
||||
tf2torch_tensor_name_prefix_tf: str = "seq2seq/decoder",
|
||||
embed_tensor_name_prefix_tf: str = None,
|
||||
):
|
||||
super().__init__(
|
||||
vocab_size=vocab_size,
|
||||
encoder_output_size=encoder_output_size,
|
||||
dropout_rate=dropout_rate,
|
||||
positional_dropout_rate=positional_dropout_rate,
|
||||
input_layer=input_layer,
|
||||
use_output_layer=use_output_layer,
|
||||
pos_enc_class=pos_enc_class,
|
||||
normalize_before=normalize_before,
|
||||
)
|
||||
if attention_dim is None:
|
||||
attention_dim = encoder_output_size
|
||||
|
||||
if input_layer == "embed":
|
||||
self.embed = torch.nn.Sequential(
|
||||
torch.nn.Embedding(vocab_size, attention_dim),
|
||||
)
|
||||
elif input_layer == "linear":
|
||||
self.embed = torch.nn.Sequential(
|
||||
torch.nn.Linear(vocab_size, attention_dim),
|
||||
torch.nn.LayerNorm(attention_dim),
|
||||
torch.nn.Dropout(dropout_rate),
|
||||
torch.nn.ReLU(),
|
||||
pos_enc_class(attention_dim, positional_dropout_rate),
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"only 'embed' or 'linear' is supported: {input_layer}")
|
||||
|
||||
self.normalize_before = normalize_before
|
||||
if self.normalize_before:
|
||||
self.after_norm = LayerNorm(attention_dim)
|
||||
if use_output_layer:
|
||||
self.output_layer = torch.nn.Linear(attention_dim, vocab_size)
|
||||
else:
|
||||
self.output_layer = None
|
||||
|
||||
self.att_layer_num = att_layer_num
|
||||
self.num_blocks = num_blocks
|
||||
if sanm_shfit is None:
|
||||
sanm_shfit = (kernel_size - 1) // 2
|
||||
self.decoders = repeat(
|
||||
att_layer_num,
|
||||
lambda lnum: DecoderLayerSANM(
|
||||
attention_dim,
|
||||
MultiHeadedAttentionSANMDecoder(
|
||||
attention_dim, self_attention_dropout_rate, kernel_size, sanm_shfit=sanm_shfit
|
||||
),
|
||||
MultiHeadedAttentionCrossAtt(
|
||||
attention_heads,
|
||||
attention_dim,
|
||||
src_attention_dropout_rate,
|
||||
encoder_output_size=encoder_output_size,
|
||||
),
|
||||
PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate),
|
||||
dropout_rate,
|
||||
normalize_before,
|
||||
concat_after,
|
||||
),
|
||||
)
|
||||
if num_blocks - att_layer_num <= 0:
|
||||
self.decoders2 = None
|
||||
else:
|
||||
self.decoders2 = repeat(
|
||||
num_blocks - att_layer_num,
|
||||
lambda lnum: DecoderLayerSANM(
|
||||
attention_dim,
|
||||
MultiHeadedAttentionSANMDecoder(
|
||||
attention_dim,
|
||||
self_attention_dropout_rate,
|
||||
kernel_size,
|
||||
sanm_shfit=sanm_shfit,
|
||||
),
|
||||
None,
|
||||
PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate),
|
||||
dropout_rate,
|
||||
normalize_before,
|
||||
concat_after,
|
||||
),
|
||||
)
|
||||
|
||||
self.decoders3 = repeat(
|
||||
1,
|
||||
lambda lnum: DecoderLayerSANM(
|
||||
attention_dim,
|
||||
None,
|
||||
None,
|
||||
PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate),
|
||||
dropout_rate,
|
||||
normalize_before,
|
||||
concat_after,
|
||||
),
|
||||
)
|
||||
if concat_embeds:
|
||||
self.embed_concat_ffn = repeat(
|
||||
1,
|
||||
lambda lnum: DecoderLayerSANM(
|
||||
attention_dim + encoder_output_size,
|
||||
None,
|
||||
None,
|
||||
PositionwiseFeedForwardDecoderSANM(
|
||||
attention_dim + encoder_output_size,
|
||||
linear_units,
|
||||
dropout_rate,
|
||||
adim=attention_dim,
|
||||
),
|
||||
dropout_rate,
|
||||
normalize_before,
|
||||
concat_after,
|
||||
),
|
||||
)
|
||||
else:
|
||||
self.embed_concat_ffn = None
|
||||
self.concat_embeds = concat_embeds
|
||||
self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch
|
||||
self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf
|
||||
self.embed_tensor_name_prefix_tf = embed_tensor_name_prefix_tf
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hs_pad: torch.Tensor,
|
||||
hlens: torch.Tensor,
|
||||
ys_in_pad: torch.Tensor,
|
||||
ys_in_lens: torch.Tensor,
|
||||
chunk_mask: torch.Tensor = None,
|
||||
pre_acoustic_embeds: torch.Tensor = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""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, )
|
||||
"""
|
||||
tgt = ys_in_pad
|
||||
tgt_mask = myutils.sequence_mask(ys_in_lens, device=tgt.device)[:, :, None]
|
||||
|
||||
memory = hs_pad
|
||||
memory_mask = myutils.sequence_mask(hlens, device=memory.device)[:, None, :]
|
||||
if chunk_mask is not None:
|
||||
memory_mask = memory_mask * chunk_mask
|
||||
if tgt_mask.size(1) != memory_mask.size(1):
|
||||
memory_mask = torch.cat((memory_mask, memory_mask[:, -2:-1, :]), dim=1)
|
||||
|
||||
x = self.embed(tgt)
|
||||
|
||||
if pre_acoustic_embeds is not None and self.concat_embeds:
|
||||
x = torch.cat((x, pre_acoustic_embeds), dim=-1)
|
||||
x, _, _, _, _ = self.embed_concat_ffn(x, None, None, None, None)
|
||||
|
||||
x, tgt_mask, memory, memory_mask, _ = self.decoders(x, tgt_mask, memory, memory_mask)
|
||||
if self.decoders2 is not None:
|
||||
x, tgt_mask, memory, memory_mask, _ = self.decoders2(x, tgt_mask, memory, memory_mask)
|
||||
x, tgt_mask, memory, memory_mask, _ = self.decoders3(x, tgt_mask, memory, memory_mask)
|
||||
if self.normalize_before:
|
||||
x = self.after_norm(x)
|
||||
if self.output_layer is not None:
|
||||
x = self.output_layer(x)
|
||||
|
||||
olens = tgt_mask.sum(1)
|
||||
return x, olens
|
||||
|
||||
def score(
|
||||
self,
|
||||
ys,
|
||||
state,
|
||||
x,
|
||||
x_mask=None,
|
||||
pre_acoustic_embeds: torch.Tensor = None,
|
||||
):
|
||||
"""Score."""
|
||||
ys_mask = myutils.sequence_mask(
|
||||
torch.tensor([len(ys)], dtype=torch.int32), device=x.device
|
||||
)[:, :, None]
|
||||
logp, state = self.forward_one_step(
|
||||
ys.unsqueeze(0),
|
||||
ys_mask,
|
||||
x.unsqueeze(0),
|
||||
memory_mask=x_mask,
|
||||
pre_acoustic_embeds=pre_acoustic_embeds,
|
||||
cache=state,
|
||||
)
|
||||
return logp.squeeze(0), state
|
||||
|
||||
def forward_one_step(
|
||||
self,
|
||||
tgt: torch.Tensor,
|
||||
tgt_mask: torch.Tensor,
|
||||
memory: torch.Tensor,
|
||||
memory_mask: torch.Tensor = None,
|
||||
pre_acoustic_embeds: torch.Tensor = None,
|
||||
cache: List[torch.Tensor] = None,
|
||||
) -> Tuple[torch.Tensor, List[torch.Tensor]]:
|
||||
"""Forward one step.
|
||||
|
||||
Args:
|
||||
tgt: input token ids, int64 (batch, maxlen_out)
|
||||
tgt_mask: input token mask, (batch, maxlen_out)
|
||||
dtype=torch.uint8 in PyTorch 1.2-
|
||||
dtype=torch.bool in PyTorch 1.2+ (include 1.2)
|
||||
memory: encoded memory, float32 (batch, maxlen_in, feat)
|
||||
cache: cached output list of (batch, max_time_out-1, size)
|
||||
Returns:
|
||||
y, cache: NN output value and cache per `self.decoders`.
|
||||
y.shape` is (batch, maxlen_out, token)
|
||||
"""
|
||||
|
||||
x = tgt[:, -1:]
|
||||
tgt_mask = None
|
||||
x = self.embed(x)
|
||||
|
||||
if pre_acoustic_embeds is not None and self.concat_embeds:
|
||||
x = torch.cat((x, pre_acoustic_embeds), dim=-1)
|
||||
x, _, _, _, _ = self.embed_concat_ffn(x, None, None, None, None)
|
||||
|
||||
if cache is None:
|
||||
cache_layer_num = len(self.decoders)
|
||||
if self.decoders2 is not None:
|
||||
cache_layer_num += len(self.decoders2)
|
||||
cache = [None] * cache_layer_num
|
||||
new_cache = []
|
||||
# for c, decoder in zip(cache, self.decoders):
|
||||
for i in range(self.att_layer_num):
|
||||
decoder = self.decoders[i]
|
||||
c = cache[i]
|
||||
x, tgt_mask, memory, memory_mask, c_ret = decoder.forward_one_step(
|
||||
x, tgt_mask, memory, memory_mask, cache=c
|
||||
)
|
||||
new_cache.append(c_ret)
|
||||
|
||||
if self.num_blocks - self.att_layer_num >= 1:
|
||||
for i in range(self.num_blocks - self.att_layer_num):
|
||||
j = i + self.att_layer_num
|
||||
decoder = self.decoders2[i]
|
||||
c = cache[j]
|
||||
x, tgt_mask, memory, memory_mask, c_ret = decoder.forward_one_step(
|
||||
x, tgt_mask, memory, memory_mask, cache=c
|
||||
)
|
||||
new_cache.append(c_ret)
|
||||
|
||||
for decoder in self.decoders3:
|
||||
x, tgt_mask, memory, memory_mask, _ = decoder.forward_one_step(
|
||||
x, tgt_mask, memory, None, cache=None
|
||||
)
|
||||
|
||||
if self.normalize_before:
|
||||
y = self.after_norm(x[:, -1])
|
||||
else:
|
||||
y = x[:, -1]
|
||||
if self.output_layer is not None:
|
||||
y = self.output_layer(y)
|
||||
y = torch.log_softmax(y, dim=-1)
|
||||
|
||||
return y, new_cache
|
||||
@@ -0,0 +1,499 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- encoding: utf-8 -*-
|
||||
# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
|
||||
# MIT License (https://opensource.org/licenses/MIT)
|
||||
|
||||
from typing import List
|
||||
from typing import Optional
|
||||
from typing import Sequence
|
||||
from typing import Tuple
|
||||
from typing import Union
|
||||
import logging
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from funasr.models.scama.chunk_utilis import overlap_chunk
|
||||
import numpy as np
|
||||
from funasr.train_utils.device_funcs import to_device
|
||||
from funasr.models.transformer.utils.nets_utils import make_pad_mask
|
||||
from funasr.models.sanm.attention import MultiHeadedAttention, MultiHeadedAttentionSANM
|
||||
from funasr.models.transformer.embedding import (
|
||||
SinusoidalPositionEncoder,
|
||||
StreamSinusoidalPositionEncoder,
|
||||
)
|
||||
from funasr.models.transformer.layer_norm import LayerNorm
|
||||
from funasr.models.transformer.utils.multi_layer_conv import Conv1dLinear
|
||||
from funasr.models.transformer.utils.multi_layer_conv import MultiLayeredConv1d
|
||||
from funasr.models.transformer.positionwise_feed_forward import (
|
||||
PositionwiseFeedForward, # noqa: H301
|
||||
)
|
||||
from funasr.models.transformer.utils.repeat import repeat
|
||||
from funasr.models.transformer.utils.subsampling import Conv2dSubsampling
|
||||
from funasr.models.transformer.utils.subsampling import Conv2dSubsampling2
|
||||
from funasr.models.transformer.utils.subsampling import Conv2dSubsampling6
|
||||
from funasr.models.transformer.utils.subsampling import Conv2dSubsampling8
|
||||
from funasr.models.transformer.utils.subsampling import TooShortUttError
|
||||
from funasr.models.transformer.utils.subsampling import check_short_utt
|
||||
from funasr.models.transformer.utils.mask import subsequent_mask, vad_mask
|
||||
|
||||
from funasr.models.ctc.ctc import CTC
|
||||
|
||||
from funasr.register import tables
|
||||
|
||||
|
||||
class EncoderLayerSANM(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_size,
|
||||
size,
|
||||
self_attn,
|
||||
feed_forward,
|
||||
dropout_rate,
|
||||
normalize_before=True,
|
||||
concat_after=False,
|
||||
stochastic_depth_rate=0.0,
|
||||
):
|
||||
"""Construct an EncoderLayer object."""
|
||||
super(EncoderLayerSANM, self).__init__()
|
||||
self.self_attn = self_attn
|
||||
self.feed_forward = feed_forward
|
||||
self.norm1 = LayerNorm(in_size)
|
||||
self.norm2 = LayerNorm(size)
|
||||
self.dropout = nn.Dropout(dropout_rate)
|
||||
self.in_size = in_size
|
||||
self.size = size
|
||||
self.normalize_before = normalize_before
|
||||
self.concat_after = concat_after
|
||||
if self.concat_after:
|
||||
self.concat_linear = nn.Linear(size + size, size)
|
||||
self.stochastic_depth_rate = stochastic_depth_rate
|
||||
self.dropout_rate = dropout_rate
|
||||
|
||||
def forward(self, x, mask, cache=None, mask_shfit_chunk=None, mask_att_chunk_encoder=None):
|
||||
"""Compute encoded features.
|
||||
|
||||
Args:
|
||||
x_input (torch.Tensor): Input tensor (#batch, time, size).
|
||||
mask (torch.Tensor): Mask tensor for the input (#batch, time).
|
||||
cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size).
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Output tensor (#batch, time, size).
|
||||
torch.Tensor: Mask tensor (#batch, time).
|
||||
|
||||
"""
|
||||
skip_layer = False
|
||||
# with stochastic depth, residual connection `x + f(x)` becomes
|
||||
# `x <- x + 1 / (1 - p) * f(x)` at training time.
|
||||
stoch_layer_coeff = 1.0
|
||||
if self.training and self.stochastic_depth_rate > 0:
|
||||
skip_layer = torch.rand(1).item() < self.stochastic_depth_rate
|
||||
stoch_layer_coeff = 1.0 / (1 - self.stochastic_depth_rate)
|
||||
|
||||
if skip_layer:
|
||||
if cache is not None:
|
||||
x = torch.cat([cache, x], dim=1)
|
||||
return x, mask
|
||||
|
||||
residual = x
|
||||
if self.normalize_before:
|
||||
x = self.norm1(x)
|
||||
|
||||
if self.concat_after:
|
||||
x_concat = torch.cat(
|
||||
(
|
||||
x,
|
||||
self.self_attn(
|
||||
x,
|
||||
mask,
|
||||
mask_shfit_chunk=mask_shfit_chunk,
|
||||
mask_att_chunk_encoder=mask_att_chunk_encoder,
|
||||
),
|
||||
),
|
||||
dim=-1,
|
||||
)
|
||||
if self.in_size == self.size:
|
||||
x = residual + stoch_layer_coeff * self.concat_linear(x_concat)
|
||||
else:
|
||||
x = stoch_layer_coeff * self.concat_linear(x_concat)
|
||||
else:
|
||||
if self.in_size == self.size:
|
||||
x = residual + stoch_layer_coeff * self.dropout(
|
||||
self.self_attn(
|
||||
x,
|
||||
mask,
|
||||
mask_shfit_chunk=mask_shfit_chunk,
|
||||
mask_att_chunk_encoder=mask_att_chunk_encoder,
|
||||
)
|
||||
)
|
||||
else:
|
||||
x = stoch_layer_coeff * self.dropout(
|
||||
self.self_attn(
|
||||
x,
|
||||
mask,
|
||||
mask_shfit_chunk=mask_shfit_chunk,
|
||||
mask_att_chunk_encoder=mask_att_chunk_encoder,
|
||||
)
|
||||
)
|
||||
if not self.normalize_before:
|
||||
x = self.norm1(x)
|
||||
|
||||
residual = x
|
||||
if self.normalize_before:
|
||||
x = self.norm2(x)
|
||||
x = residual + stoch_layer_coeff * self.dropout(self.feed_forward(x))
|
||||
if not self.normalize_before:
|
||||
x = self.norm2(x)
|
||||
|
||||
return x, mask, cache, mask_shfit_chunk, mask_att_chunk_encoder
|
||||
|
||||
def forward_chunk(self, x, cache=None, chunk_size=None, look_back=0):
|
||||
"""Compute encoded features.
|
||||
|
||||
Args:
|
||||
x_input (torch.Tensor): Input tensor (#batch, time, size).
|
||||
mask (torch.Tensor): Mask tensor for the input (#batch, time).
|
||||
cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size).
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Output tensor (#batch, time, size).
|
||||
torch.Tensor: Mask tensor (#batch, time).
|
||||
|
||||
"""
|
||||
|
||||
residual = x
|
||||
if self.normalize_before:
|
||||
x = self.norm1(x)
|
||||
|
||||
if self.in_size == self.size:
|
||||
attn, cache = self.self_attn.forward_chunk(x, cache, chunk_size, look_back)
|
||||
x = residual + attn
|
||||
else:
|
||||
x, cache = self.self_attn.forward_chunk(x, cache, chunk_size, look_back)
|
||||
|
||||
if not self.normalize_before:
|
||||
x = self.norm1(x)
|
||||
|
||||
residual = x
|
||||
if self.normalize_before:
|
||||
x = self.norm2(x)
|
||||
x = residual + self.feed_forward(x)
|
||||
if not self.normalize_before:
|
||||
x = self.norm2(x)
|
||||
|
||||
return x, cache
|
||||
|
||||
|
||||
@tables.register("encoder_classes", "SANMEncoderChunkOpt")
|
||||
class SANMEncoderChunkOpt(nn.Module):
|
||||
"""
|
||||
Author: Shiliang Zhang, Zhifu Gao, Haoneng Luo, Ming Lei, Jie Gao, Zhijie Yan, Lei Xie
|
||||
SCAMA: Streaming chunk-aware multihead attention for online end-to-end speech recognition
|
||||
https://arxiv.org/abs/2006.01712
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_size: int,
|
||||
output_size: int = 256,
|
||||
attention_heads: int = 4,
|
||||
linear_units: int = 2048,
|
||||
num_blocks: int = 6,
|
||||
dropout_rate: float = 0.1,
|
||||
positional_dropout_rate: float = 0.1,
|
||||
attention_dropout_rate: float = 0.0,
|
||||
input_layer: Optional[str] = "conv2d",
|
||||
pos_enc_class=SinusoidalPositionEncoder,
|
||||
normalize_before: bool = True,
|
||||
concat_after: bool = False,
|
||||
positionwise_layer_type: str = "linear",
|
||||
positionwise_conv_kernel_size: int = 1,
|
||||
padding_idx: int = -1,
|
||||
interctc_layer_idx: List[int] = [],
|
||||
interctc_use_conditioning: bool = False,
|
||||
kernel_size: int = 11,
|
||||
sanm_shfit: int = 0,
|
||||
selfattention_layer_type: str = "sanm",
|
||||
chunk_size: Union[int, Sequence[int]] = (16,),
|
||||
stride: Union[int, Sequence[int]] = (10,),
|
||||
pad_left: Union[int, Sequence[int]] = (0,),
|
||||
encoder_att_look_back_factor: Union[int, Sequence[int]] = (1,),
|
||||
decoder_att_look_back_factor: Union[int, Sequence[int]] = (1,),
|
||||
tf2torch_tensor_name_prefix_torch: str = "encoder",
|
||||
tf2torch_tensor_name_prefix_tf: str = "seq2seq/encoder",
|
||||
):
|
||||
super().__init__()
|
||||
self._output_size = output_size
|
||||
|
||||
if input_layer == "linear":
|
||||
self.embed = torch.nn.Sequential(
|
||||
torch.nn.Linear(input_size, output_size),
|
||||
torch.nn.LayerNorm(output_size),
|
||||
torch.nn.Dropout(dropout_rate),
|
||||
torch.nn.ReLU(),
|
||||
pos_enc_class(output_size, positional_dropout_rate),
|
||||
)
|
||||
elif input_layer == "conv2d":
|
||||
self.embed = Conv2dSubsampling(input_size, output_size, dropout_rate)
|
||||
elif input_layer == "conv2d2":
|
||||
self.embed = Conv2dSubsampling2(input_size, output_size, dropout_rate)
|
||||
elif input_layer == "conv2d6":
|
||||
self.embed = Conv2dSubsampling6(input_size, output_size, dropout_rate)
|
||||
elif input_layer == "conv2d8":
|
||||
self.embed = Conv2dSubsampling8(input_size, output_size, dropout_rate)
|
||||
elif input_layer == "embed":
|
||||
self.embed = torch.nn.Sequential(
|
||||
torch.nn.Embedding(input_size, output_size, padding_idx=padding_idx),
|
||||
pos_enc_class(output_size, positional_dropout_rate),
|
||||
)
|
||||
elif input_layer is None:
|
||||
if input_size == output_size:
|
||||
self.embed = None
|
||||
else:
|
||||
self.embed = torch.nn.Linear(input_size, output_size)
|
||||
elif input_layer == "pe":
|
||||
self.embed = SinusoidalPositionEncoder()
|
||||
elif input_layer == "pe_online":
|
||||
self.embed = StreamSinusoidalPositionEncoder()
|
||||
else:
|
||||
raise ValueError("unknown input_layer: " + input_layer)
|
||||
self.normalize_before = normalize_before
|
||||
if positionwise_layer_type == "linear":
|
||||
positionwise_layer = PositionwiseFeedForward
|
||||
positionwise_layer_args = (
|
||||
output_size,
|
||||
linear_units,
|
||||
dropout_rate,
|
||||
)
|
||||
elif positionwise_layer_type == "conv1d":
|
||||
positionwise_layer = MultiLayeredConv1d
|
||||
positionwise_layer_args = (
|
||||
output_size,
|
||||
linear_units,
|
||||
positionwise_conv_kernel_size,
|
||||
dropout_rate,
|
||||
)
|
||||
elif positionwise_layer_type == "conv1d-linear":
|
||||
positionwise_layer = Conv1dLinear
|
||||
positionwise_layer_args = (
|
||||
output_size,
|
||||
linear_units,
|
||||
positionwise_conv_kernel_size,
|
||||
dropout_rate,
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError("Support only linear or conv1d.")
|
||||
|
||||
if selfattention_layer_type == "selfattn":
|
||||
encoder_selfattn_layer = MultiHeadedAttention
|
||||
encoder_selfattn_layer_args = (
|
||||
attention_heads,
|
||||
output_size,
|
||||
attention_dropout_rate,
|
||||
)
|
||||
elif selfattention_layer_type == "sanm":
|
||||
encoder_selfattn_layer = MultiHeadedAttentionSANM
|
||||
encoder_selfattn_layer_args0 = (
|
||||
attention_heads,
|
||||
input_size,
|
||||
output_size,
|
||||
attention_dropout_rate,
|
||||
kernel_size,
|
||||
sanm_shfit,
|
||||
)
|
||||
|
||||
encoder_selfattn_layer_args = (
|
||||
attention_heads,
|
||||
output_size,
|
||||
output_size,
|
||||
attention_dropout_rate,
|
||||
kernel_size,
|
||||
sanm_shfit,
|
||||
)
|
||||
self.encoders0 = repeat(
|
||||
1,
|
||||
lambda lnum: EncoderLayerSANM(
|
||||
input_size,
|
||||
output_size,
|
||||
encoder_selfattn_layer(*encoder_selfattn_layer_args0),
|
||||
positionwise_layer(*positionwise_layer_args),
|
||||
dropout_rate,
|
||||
normalize_before,
|
||||
concat_after,
|
||||
),
|
||||
)
|
||||
|
||||
self.encoders = repeat(
|
||||
num_blocks - 1,
|
||||
lambda lnum: EncoderLayerSANM(
|
||||
output_size,
|
||||
output_size,
|
||||
encoder_selfattn_layer(*encoder_selfattn_layer_args),
|
||||
positionwise_layer(*positionwise_layer_args),
|
||||
dropout_rate,
|
||||
normalize_before,
|
||||
concat_after,
|
||||
),
|
||||
)
|
||||
if self.normalize_before:
|
||||
self.after_norm = LayerNorm(output_size)
|
||||
|
||||
self.interctc_layer_idx = interctc_layer_idx
|
||||
if len(interctc_layer_idx) > 0:
|
||||
assert 0 < min(interctc_layer_idx) and max(interctc_layer_idx) < num_blocks
|
||||
self.interctc_use_conditioning = interctc_use_conditioning
|
||||
self.conditioning_layer = None
|
||||
shfit_fsmn = (kernel_size - 1) // 2
|
||||
self.overlap_chunk_cls = overlap_chunk(
|
||||
chunk_size=chunk_size,
|
||||
stride=stride,
|
||||
pad_left=pad_left,
|
||||
shfit_fsmn=shfit_fsmn,
|
||||
encoder_att_look_back_factor=encoder_att_look_back_factor,
|
||||
decoder_att_look_back_factor=decoder_att_look_back_factor,
|
||||
)
|
||||
self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch
|
||||
self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf
|
||||
|
||||
def output_size(self) -> int:
|
||||
return self._output_size
|
||||
|
||||
def forward(
|
||||
self,
|
||||
xs_pad: torch.Tensor,
|
||||
ilens: torch.Tensor,
|
||||
prev_states: torch.Tensor = None,
|
||||
ctc: CTC = None,
|
||||
ind: int = 0,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
|
||||
"""Embed positions in tensor.
|
||||
|
||||
Args:
|
||||
xs_pad: input tensor (B, L, D)
|
||||
ilens: input length (B)
|
||||
prev_states: Not to be used now.
|
||||
Returns:
|
||||
position embedded tensor and mask
|
||||
"""
|
||||
masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device)
|
||||
xs_pad *= self.output_size() ** 0.5
|
||||
if self.embed is None:
|
||||
xs_pad = xs_pad
|
||||
elif (
|
||||
isinstance(self.embed, Conv2dSubsampling)
|
||||
or isinstance(self.embed, Conv2dSubsampling2)
|
||||
or isinstance(self.embed, Conv2dSubsampling6)
|
||||
or isinstance(self.embed, Conv2dSubsampling8)
|
||||
):
|
||||
short_status, limit_size = check_short_utt(self.embed, xs_pad.size(1))
|
||||
if short_status:
|
||||
raise TooShortUttError(
|
||||
f"has {xs_pad.size(1)} frames and is too short for subsampling "
|
||||
+ f"(it needs more than {limit_size} frames), return empty results",
|
||||
xs_pad.size(1),
|
||||
limit_size,
|
||||
)
|
||||
xs_pad, masks = self.embed(xs_pad, masks)
|
||||
else:
|
||||
xs_pad = self.embed(xs_pad)
|
||||
|
||||
mask_shfit_chunk, mask_att_chunk_encoder = None, None
|
||||
if self.overlap_chunk_cls is not None:
|
||||
ilens = masks.squeeze(1).sum(1)
|
||||
chunk_outs = self.overlap_chunk_cls.gen_chunk_mask(ilens, ind)
|
||||
xs_pad, ilens = self.overlap_chunk_cls.split_chunk(xs_pad, ilens, chunk_outs=chunk_outs)
|
||||
masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device)
|
||||
mask_shfit_chunk = self.overlap_chunk_cls.get_mask_shfit_chunk(
|
||||
chunk_outs, xs_pad.device, xs_pad.size(0), dtype=xs_pad.dtype
|
||||
)
|
||||
mask_att_chunk_encoder = self.overlap_chunk_cls.get_mask_att_chunk_encoder(
|
||||
chunk_outs, xs_pad.device, xs_pad.size(0), dtype=xs_pad.dtype
|
||||
)
|
||||
|
||||
encoder_outs = self.encoders0(xs_pad, masks, None, mask_shfit_chunk, mask_att_chunk_encoder)
|
||||
xs_pad, masks = encoder_outs[0], encoder_outs[1]
|
||||
intermediate_outs = []
|
||||
if len(self.interctc_layer_idx) == 0:
|
||||
encoder_outs = self.encoders(
|
||||
xs_pad, masks, None, mask_shfit_chunk, mask_att_chunk_encoder
|
||||
)
|
||||
xs_pad, masks = encoder_outs[0], encoder_outs[1]
|
||||
else:
|
||||
for layer_idx, encoder_layer in enumerate(self.encoders):
|
||||
encoder_outs = encoder_layer(
|
||||
xs_pad, masks, None, mask_shfit_chunk, mask_att_chunk_encoder
|
||||
)
|
||||
xs_pad, masks = encoder_outs[0], encoder_outs[1]
|
||||
if layer_idx + 1 in self.interctc_layer_idx:
|
||||
encoder_out = xs_pad
|
||||
|
||||
# intermediate outputs are also normalized
|
||||
if self.normalize_before:
|
||||
encoder_out = self.after_norm(encoder_out)
|
||||
|
||||
intermediate_outs.append((layer_idx + 1, encoder_out))
|
||||
|
||||
if self.interctc_use_conditioning:
|
||||
ctc_out = ctc.softmax(encoder_out)
|
||||
xs_pad = xs_pad + self.conditioning_layer(ctc_out)
|
||||
|
||||
if self.normalize_before:
|
||||
xs_pad = self.after_norm(xs_pad)
|
||||
|
||||
olens = masks.squeeze(1).sum(1)
|
||||
if len(intermediate_outs) > 0:
|
||||
return (xs_pad, intermediate_outs), olens, None
|
||||
return xs_pad, olens, None
|
||||
|
||||
def _add_overlap_chunk(self, feats: np.ndarray, cache: dict = {}):
|
||||
if len(cache) == 0:
|
||||
return feats
|
||||
cache["feats"] = to_device(cache["feats"], device=feats.device)
|
||||
overlap_feats = torch.cat((cache["feats"], feats), dim=1)
|
||||
cache["feats"] = overlap_feats[:, -(cache["chunk_size"][0] + cache["chunk_size"][2]) :, :]
|
||||
return overlap_feats
|
||||
|
||||
def forward_chunk(
|
||||
self,
|
||||
xs_pad: torch.Tensor,
|
||||
ilens: torch.Tensor,
|
||||
cache: dict = None,
|
||||
**kwargs,
|
||||
):
|
||||
is_final = kwargs.get("is_final", False)
|
||||
xs_pad *= self.output_size() ** 0.5
|
||||
if self.embed is None:
|
||||
xs_pad = xs_pad
|
||||
else:
|
||||
xs_pad = self.embed(xs_pad, cache)
|
||||
if cache["tail_chunk"]:
|
||||
xs_pad = to_device(cache["feats"], device=xs_pad.device)
|
||||
else:
|
||||
xs_pad = self._add_overlap_chunk(xs_pad, cache)
|
||||
if cache["opt"] is None:
|
||||
cache_layer_num = len(self.encoders0) + len(self.encoders)
|
||||
new_cache = [None] * cache_layer_num
|
||||
else:
|
||||
new_cache = cache["opt"]
|
||||
|
||||
for layer_idx, encoder_layer in enumerate(self.encoders0):
|
||||
encoder_outs = encoder_layer.forward_chunk(
|
||||
xs_pad, new_cache[layer_idx], cache["chunk_size"], cache["encoder_chunk_look_back"]
|
||||
)
|
||||
xs_pad, new_cache[0] = encoder_outs[0], encoder_outs[1]
|
||||
|
||||
for layer_idx, encoder_layer in enumerate(self.encoders):
|
||||
encoder_outs = encoder_layer.forward_chunk(
|
||||
xs_pad,
|
||||
new_cache[layer_idx + len(self.encoders0)],
|
||||
cache["chunk_size"],
|
||||
cache["encoder_chunk_look_back"],
|
||||
)
|
||||
xs_pad, new_cache[layer_idx + len(self.encoders0)] = encoder_outs[0], encoder_outs[1]
|
||||
|
||||
if self.normalize_before:
|
||||
xs_pad = self.after_norm(xs_pad)
|
||||
if cache["encoder_chunk_look_back"] > 0 or cache["encoder_chunk_look_back"] == -1:
|
||||
cache["opt"] = new_cache
|
||||
|
||||
return xs_pad, ilens, None
|
||||
@@ -0,0 +1,739 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- encoding: utf-8 -*-
|
||||
# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
|
||||
# MIT License (https://opensource.org/licenses/MIT)
|
||||
|
||||
import time
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.functional as F
|
||||
import logging
|
||||
from typing import Dict, Tuple
|
||||
from contextlib import contextmanager
|
||||
from distutils.version import LooseVersion
|
||||
|
||||
from funasr.register import tables
|
||||
from funasr.models.ctc.ctc import CTC
|
||||
from funasr.utils import postprocess_utils
|
||||
from funasr.metrics.compute_acc import th_accuracy
|
||||
from funasr.utils.datadir_writer import DatadirWriter
|
||||
from funasr.models.paraformer.model import Paraformer
|
||||
from funasr.models.paraformer.search import Hypothesis
|
||||
from funasr.models.paraformer.cif_predictor import mae_loss
|
||||
from funasr.train_utils.device_funcs import force_gatherable
|
||||
from funasr.losses.label_smoothing_loss import LabelSmoothingLoss
|
||||
from funasr.models.transformer.utils.add_sos_eos import add_sos_eos
|
||||
from funasr.models.transformer.utils.nets_utils import make_pad_mask, pad_list
|
||||
from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
|
||||
from funasr.models.scama.utils import sequence_mask
|
||||
|
||||
if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
|
||||
from torch.cuda.amp import autocast
|
||||
else:
|
||||
# Nothing to do if torch<1.6.0
|
||||
@contextmanager
|
||||
def autocast(enabled=True):
|
||||
yield
|
||||
|
||||
|
||||
@tables.register("model_classes", "SCAMA")
|
||||
class SCAMA(nn.Module):
|
||||
"""
|
||||
Author: Shiliang Zhang, Zhifu Gao, Haoneng Luo, Ming Lei, Jie Gao, Zhijie Yan, Lei Xie
|
||||
SCAMA: Streaming chunk-aware multihead attention for online end-to-end speech recognition
|
||||
https://arxiv.org/abs/2006.01712
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
specaug: str = None,
|
||||
specaug_conf: dict = None,
|
||||
normalize: str = None,
|
||||
normalize_conf: dict = None,
|
||||
encoder: str = None,
|
||||
encoder_conf: dict = None,
|
||||
decoder: str = None,
|
||||
decoder_conf: dict = None,
|
||||
ctc: str = None,
|
||||
ctc_conf: dict = None,
|
||||
ctc_weight: float = 0.5,
|
||||
predictor: str = None,
|
||||
predictor_conf: dict = None,
|
||||
predictor_bias: int = 0,
|
||||
predictor_weight: float = 0.0,
|
||||
input_size: int = 80,
|
||||
vocab_size: int = -1,
|
||||
ignore_id: int = -1,
|
||||
blank_id: int = 0,
|
||||
sos: int = 1,
|
||||
eos: int = 2,
|
||||
lsm_weight: float = 0.0,
|
||||
length_normalized_loss: bool = False,
|
||||
share_embedding: bool = False,
|
||||
**kwargs,
|
||||
):
|
||||
|
||||
super().__init__()
|
||||
|
||||
if specaug is not None:
|
||||
specaug_class = tables.specaug_classes.get(specaug)
|
||||
specaug = specaug_class(**specaug_conf)
|
||||
|
||||
if normalize is not None:
|
||||
normalize_class = tables.normalize_classes.get(normalize)
|
||||
normalize = normalize_class(**normalize_conf)
|
||||
|
||||
encoder_class = tables.encoder_classes.get(encoder)
|
||||
encoder = encoder_class(input_size=input_size, **encoder_conf)
|
||||
encoder_output_size = encoder.output_size()
|
||||
|
||||
decoder_class = tables.decoder_classes.get(decoder)
|
||||
decoder = decoder_class(
|
||||
vocab_size=vocab_size,
|
||||
encoder_output_size=encoder_output_size,
|
||||
**decoder_conf,
|
||||
)
|
||||
if ctc_weight > 0.0:
|
||||
|
||||
if ctc_conf is None:
|
||||
ctc_conf = {}
|
||||
|
||||
ctc = CTC(odim=vocab_size, encoder_output_size=encoder_output_size, **ctc_conf)
|
||||
|
||||
predictor_class = tables.predictor_classes.get(predictor)
|
||||
predictor = predictor_class(**predictor_conf)
|
||||
|
||||
# note that eos is the same as sos (equivalent ID)
|
||||
self.blank_id = blank_id
|
||||
self.sos = sos if sos is not None else vocab_size - 1
|
||||
self.eos = eos if eos is not None else vocab_size - 1
|
||||
self.vocab_size = vocab_size
|
||||
self.ignore_id = ignore_id
|
||||
self.ctc_weight = ctc_weight
|
||||
|
||||
self.specaug = specaug
|
||||
self.normalize = normalize
|
||||
|
||||
self.encoder = encoder
|
||||
|
||||
if ctc_weight == 1.0:
|
||||
self.decoder = None
|
||||
else:
|
||||
self.decoder = decoder
|
||||
|
||||
self.criterion_att = LabelSmoothingLoss(
|
||||
size=vocab_size,
|
||||
padding_idx=ignore_id,
|
||||
smoothing=lsm_weight,
|
||||
normalize_length=length_normalized_loss,
|
||||
)
|
||||
|
||||
if ctc_weight == 0.0:
|
||||
self.ctc = None
|
||||
else:
|
||||
self.ctc = ctc
|
||||
|
||||
self.predictor = predictor
|
||||
self.predictor_weight = predictor_weight
|
||||
self.predictor_bias = predictor_bias
|
||||
|
||||
self.criterion_pre = mae_loss(normalize_length=length_normalized_loss)
|
||||
|
||||
self.share_embedding = share_embedding
|
||||
if self.share_embedding:
|
||||
self.decoder.embed = None
|
||||
|
||||
self.length_normalized_loss = length_normalized_loss
|
||||
self.beam_search = None
|
||||
self.error_calculator = None
|
||||
|
||||
if self.encoder.overlap_chunk_cls is not None:
|
||||
from funasr.models.scama.chunk_utilis import (
|
||||
build_scama_mask_for_cross_attention_decoder,
|
||||
)
|
||||
|
||||
self.build_scama_mask_for_cross_attention_decoder_fn = (
|
||||
build_scama_mask_for_cross_attention_decoder
|
||||
)
|
||||
self.decoder_attention_chunk_type = kwargs.get("decoder_attention_chunk_type", "chunk")
|
||||
|
||||
def forward(
|
||||
self,
|
||||
speech: torch.Tensor,
|
||||
speech_lengths: torch.Tensor,
|
||||
text: torch.Tensor,
|
||||
text_lengths: torch.Tensor,
|
||||
**kwargs,
|
||||
) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]:
|
||||
"""Encoder + Decoder + Calc loss
|
||||
Args:
|
||||
speech: (Batch, Length, ...)
|
||||
speech_lengths: (Batch, )
|
||||
text: (Batch, Length)
|
||||
text_lengths: (Batch,)
|
||||
"""
|
||||
|
||||
decoding_ind = kwargs.get("decoding_ind")
|
||||
if len(text_lengths.size()) > 1:
|
||||
text_lengths = text_lengths[:, 0]
|
||||
if len(speech_lengths.size()) > 1:
|
||||
speech_lengths = speech_lengths[:, 0]
|
||||
|
||||
batch_size = speech.shape[0]
|
||||
|
||||
# Encoder
|
||||
ind = self.encoder.overlap_chunk_cls.random_choice(self.training, decoding_ind)
|
||||
encoder_out, encoder_out_lens = self.encode(speech, speech_lengths, ind=ind)
|
||||
|
||||
loss_ctc, cer_ctc = None, None
|
||||
loss_pre = None
|
||||
stats = dict()
|
||||
|
||||
# decoder: CTC branch
|
||||
|
||||
if self.ctc_weight > 0.0:
|
||||
|
||||
encoder_out_ctc, encoder_out_lens_ctc = self.encoder.overlap_chunk_cls.remove_chunk(
|
||||
encoder_out, encoder_out_lens, chunk_outs=None
|
||||
)
|
||||
|
||||
loss_ctc, cer_ctc = self._calc_ctc_loss(
|
||||
encoder_out_ctc, encoder_out_lens_ctc, text, text_lengths
|
||||
)
|
||||
# Collect CTC branch stats
|
||||
stats["loss_ctc"] = loss_ctc.detach() if loss_ctc is not None else None
|
||||
stats["cer_ctc"] = cer_ctc
|
||||
|
||||
# decoder: Attention decoder branch
|
||||
loss_att, acc_att, cer_att, wer_att, loss_pre = self._calc_att_predictor_loss(
|
||||
encoder_out, encoder_out_lens, text, text_lengths
|
||||
)
|
||||
|
||||
# 3. CTC-Att loss definition
|
||||
if self.ctc_weight == 0.0:
|
||||
loss = loss_att + loss_pre * self.predictor_weight
|
||||
else:
|
||||
loss = (
|
||||
self.ctc_weight * loss_ctc
|
||||
+ (1 - self.ctc_weight) * loss_att
|
||||
+ loss_pre * self.predictor_weight
|
||||
)
|
||||
|
||||
# Collect Attn branch stats
|
||||
stats["loss_att"] = loss_att.detach() if loss_att is not None else None
|
||||
stats["acc"] = acc_att
|
||||
stats["cer"] = cer_att
|
||||
stats["wer"] = wer_att
|
||||
stats["loss_pre"] = loss_pre.detach().cpu() if loss_pre is not None else None
|
||||
|
||||
stats["loss"] = torch.clone(loss.detach())
|
||||
|
||||
# force_gatherable: to-device and to-tensor if scalar for DataParallel
|
||||
if self.length_normalized_loss:
|
||||
batch_size = (text_lengths + self.predictor_bias).sum()
|
||||
loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device)
|
||||
return loss, stats, weight
|
||||
|
||||
def encode(
|
||||
self,
|
||||
speech: torch.Tensor,
|
||||
speech_lengths: torch.Tensor,
|
||||
**kwargs,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Encoder. Note that this method is used by asr_inference.py
|
||||
Args:
|
||||
speech: (Batch, Length, ...)
|
||||
speech_lengths: (Batch, )
|
||||
ind: int
|
||||
"""
|
||||
with autocast(False):
|
||||
|
||||
# Data augmentation
|
||||
if self.specaug is not None and self.training:
|
||||
speech, speech_lengths = self.specaug(speech, speech_lengths)
|
||||
|
||||
# Normalization for feature: e.g. Global-CMVN, Utterance-CMVN
|
||||
if self.normalize is not None:
|
||||
speech, speech_lengths = self.normalize(speech, speech_lengths)
|
||||
|
||||
# Forward encoder
|
||||
encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths)
|
||||
if isinstance(encoder_out, tuple):
|
||||
encoder_out = encoder_out[0]
|
||||
|
||||
return encoder_out, encoder_out_lens
|
||||
|
||||
def encode_chunk(
|
||||
self,
|
||||
speech: torch.Tensor,
|
||||
speech_lengths: torch.Tensor,
|
||||
cache: dict = None,
|
||||
**kwargs,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Frontend + Encoder. Note that this method is used by asr_inference.py
|
||||
Args:
|
||||
speech: (Batch, Length, ...)
|
||||
speech_lengths: (Batch, )
|
||||
ind: int
|
||||
"""
|
||||
with autocast(False):
|
||||
|
||||
# Data augmentation
|
||||
if self.specaug is not None and self.training:
|
||||
speech, speech_lengths = self.specaug(speech, speech_lengths)
|
||||
|
||||
# Normalization for feature: e.g. Global-CMVN, Utterance-CMVN
|
||||
if self.normalize is not None:
|
||||
speech, speech_lengths = self.normalize(speech, speech_lengths)
|
||||
|
||||
# Forward encoder
|
||||
encoder_out, encoder_out_lens, _ = self.encoder.forward_chunk(
|
||||
speech, speech_lengths, cache=cache["encoder"]
|
||||
)
|
||||
if isinstance(encoder_out, tuple):
|
||||
encoder_out = encoder_out[0]
|
||||
|
||||
return encoder_out, torch.tensor([encoder_out.size(1)])
|
||||
|
||||
def calc_predictor_chunk(self, encoder_out, encoder_out_lens, cache=None, **kwargs):
|
||||
is_final = kwargs.get("is_final", False)
|
||||
|
||||
return self.predictor.forward_chunk(encoder_out, cache["encoder"], is_final=is_final)
|
||||
|
||||
def _calc_att_predictor_loss(
|
||||
self,
|
||||
encoder_out: torch.Tensor,
|
||||
encoder_out_lens: torch.Tensor,
|
||||
ys_pad: torch.Tensor,
|
||||
ys_pad_lens: torch.Tensor,
|
||||
):
|
||||
ys_in_pad, ys_out_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id)
|
||||
ys_in_lens = ys_pad_lens + 1
|
||||
|
||||
encoder_out_mask = sequence_mask(
|
||||
encoder_out_lens,
|
||||
maxlen=encoder_out.size(1),
|
||||
dtype=encoder_out.dtype,
|
||||
device=encoder_out.device,
|
||||
)[:, None, :]
|
||||
mask_chunk_predictor = None
|
||||
if self.encoder.overlap_chunk_cls is not None:
|
||||
mask_chunk_predictor = self.encoder.overlap_chunk_cls.get_mask_chunk_predictor(
|
||||
None, device=encoder_out.device, batch_size=encoder_out.size(0)
|
||||
)
|
||||
mask_shfit_chunk = self.encoder.overlap_chunk_cls.get_mask_shfit_chunk(
|
||||
None, device=encoder_out.device, batch_size=encoder_out.size(0)
|
||||
)
|
||||
encoder_out = encoder_out * mask_shfit_chunk
|
||||
pre_acoustic_embeds, pre_token_length, pre_alphas, _ = self.predictor(
|
||||
encoder_out,
|
||||
ys_out_pad,
|
||||
encoder_out_mask,
|
||||
ignore_id=self.ignore_id,
|
||||
mask_chunk_predictor=mask_chunk_predictor,
|
||||
target_label_length=ys_in_lens,
|
||||
)
|
||||
predictor_alignments, predictor_alignments_len = self.predictor.gen_frame_alignments(
|
||||
pre_alphas, encoder_out_lens
|
||||
)
|
||||
|
||||
encoder_chunk_size = self.encoder.overlap_chunk_cls.chunk_size_pad_shift_cur
|
||||
attention_chunk_center_bias = 0
|
||||
attention_chunk_size = encoder_chunk_size
|
||||
decoder_att_look_back_factor = (
|
||||
self.encoder.overlap_chunk_cls.decoder_att_look_back_factor_cur
|
||||
)
|
||||
mask_shift_att_chunk_decoder = (
|
||||
self.encoder.overlap_chunk_cls.get_mask_shift_att_chunk_decoder(
|
||||
None, device=encoder_out.device, batch_size=encoder_out.size(0)
|
||||
)
|
||||
)
|
||||
scama_mask = self.build_scama_mask_for_cross_attention_decoder_fn(
|
||||
predictor_alignments=predictor_alignments,
|
||||
encoder_sequence_length=encoder_out_lens,
|
||||
chunk_size=1,
|
||||
encoder_chunk_size=encoder_chunk_size,
|
||||
attention_chunk_center_bias=attention_chunk_center_bias,
|
||||
attention_chunk_size=attention_chunk_size,
|
||||
attention_chunk_type=self.decoder_attention_chunk_type,
|
||||
step=None,
|
||||
predictor_mask_chunk_hopping=mask_chunk_predictor,
|
||||
decoder_att_look_back_factor=decoder_att_look_back_factor,
|
||||
mask_shift_att_chunk_decoder=mask_shift_att_chunk_decoder,
|
||||
target_length=ys_in_lens,
|
||||
is_training=self.training,
|
||||
)
|
||||
|
||||
# try:
|
||||
# 1. Forward decoder
|
||||
decoder_out, _ = self.decoder(
|
||||
encoder_out,
|
||||
encoder_out_lens,
|
||||
ys_in_pad,
|
||||
ys_in_lens,
|
||||
chunk_mask=scama_mask,
|
||||
pre_acoustic_embeds=pre_acoustic_embeds,
|
||||
)
|
||||
|
||||
# 2. Compute attention loss
|
||||
loss_att = self.criterion_att(decoder_out, ys_out_pad)
|
||||
acc_att = th_accuracy(
|
||||
decoder_out.view(-1, self.vocab_size),
|
||||
ys_out_pad,
|
||||
ignore_label=self.ignore_id,
|
||||
)
|
||||
# predictor loss
|
||||
loss_pre = self.criterion_pre(ys_in_lens.type_as(pre_token_length), pre_token_length)
|
||||
# Compute cer/wer using attention-decoder
|
||||
if self.training or self.error_calculator is None:
|
||||
cer_att, wer_att = None, None
|
||||
else:
|
||||
ys_hat = decoder_out.argmax(dim=-1)
|
||||
cer_att, wer_att = self.error_calculator(ys_hat.cpu(), ys_pad.cpu())
|
||||
|
||||
return loss_att, acc_att, cer_att, wer_att, loss_pre
|
||||
|
||||
def calc_predictor_mask(
|
||||
self,
|
||||
encoder_out: torch.Tensor,
|
||||
encoder_out_lens: torch.Tensor,
|
||||
ys_pad: torch.Tensor = None,
|
||||
ys_pad_lens: torch.Tensor = None,
|
||||
):
|
||||
# ys_in_pad, ys_out_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id)
|
||||
# ys_in_lens = ys_pad_lens + 1
|
||||
ys_out_pad, ys_in_lens = None, None
|
||||
|
||||
encoder_out_mask = sequence_mask(
|
||||
encoder_out_lens,
|
||||
maxlen=encoder_out.size(1),
|
||||
dtype=encoder_out.dtype,
|
||||
device=encoder_out.device,
|
||||
)[:, None, :]
|
||||
mask_chunk_predictor = None
|
||||
|
||||
mask_chunk_predictor = self.encoder.overlap_chunk_cls.get_mask_chunk_predictor(
|
||||
None, device=encoder_out.device, batch_size=encoder_out.size(0)
|
||||
)
|
||||
mask_shfit_chunk = self.encoder.overlap_chunk_cls.get_mask_shfit_chunk(
|
||||
None, device=encoder_out.device, batch_size=encoder_out.size(0)
|
||||
)
|
||||
encoder_out = encoder_out * mask_shfit_chunk
|
||||
pre_acoustic_embeds, pre_token_length, pre_alphas, _ = self.predictor(
|
||||
encoder_out,
|
||||
ys_out_pad,
|
||||
encoder_out_mask,
|
||||
ignore_id=self.ignore_id,
|
||||
mask_chunk_predictor=mask_chunk_predictor,
|
||||
target_label_length=ys_in_lens,
|
||||
)
|
||||
predictor_alignments, predictor_alignments_len = self.predictor.gen_frame_alignments(
|
||||
pre_alphas, encoder_out_lens
|
||||
)
|
||||
|
||||
encoder_chunk_size = self.encoder.overlap_chunk_cls.chunk_size_pad_shift_cur
|
||||
attention_chunk_center_bias = 0
|
||||
attention_chunk_size = encoder_chunk_size
|
||||
decoder_att_look_back_factor = (
|
||||
self.encoder.overlap_chunk_cls.decoder_att_look_back_factor_cur
|
||||
)
|
||||
mask_shift_att_chunk_decoder = (
|
||||
self.encoder.overlap_chunk_cls.get_mask_shift_att_chunk_decoder(
|
||||
None, device=encoder_out.device, batch_size=encoder_out.size(0)
|
||||
)
|
||||
)
|
||||
scama_mask = self.build_scama_mask_for_cross_attention_decoder_fn(
|
||||
predictor_alignments=predictor_alignments,
|
||||
encoder_sequence_length=encoder_out_lens,
|
||||
chunk_size=1,
|
||||
encoder_chunk_size=encoder_chunk_size,
|
||||
attention_chunk_center_bias=attention_chunk_center_bias,
|
||||
attention_chunk_size=attention_chunk_size,
|
||||
attention_chunk_type=self.decoder_attention_chunk_type,
|
||||
step=None,
|
||||
predictor_mask_chunk_hopping=mask_chunk_predictor,
|
||||
decoder_att_look_back_factor=decoder_att_look_back_factor,
|
||||
mask_shift_att_chunk_decoder=mask_shift_att_chunk_decoder,
|
||||
target_length=ys_in_lens,
|
||||
is_training=self.training,
|
||||
)
|
||||
|
||||
return (
|
||||
pre_acoustic_embeds,
|
||||
pre_token_length,
|
||||
predictor_alignments,
|
||||
predictor_alignments_len,
|
||||
scama_mask,
|
||||
)
|
||||
|
||||
def init_beam_search(
|
||||
self,
|
||||
**kwargs,
|
||||
):
|
||||
|
||||
from funasr.models.scama.beam_search import BeamSearchScamaStreaming
|
||||
|
||||
from funasr.models.transformer.scorers.ctc import CTCPrefixScorer
|
||||
from funasr.models.transformer.scorers.length_bonus import LengthBonus
|
||||
|
||||
# 1. Build ASR model
|
||||
scorers = {}
|
||||
|
||||
if self.ctc != None:
|
||||
ctc = CTCPrefixScorer(ctc=self.ctc, eos=self.eos)
|
||||
scorers.update(ctc=ctc)
|
||||
token_list = kwargs.get("token_list")
|
||||
scorers.update(
|
||||
decoder=self.decoder,
|
||||
length_bonus=LengthBonus(len(token_list)),
|
||||
)
|
||||
|
||||
# 3. Build ngram model
|
||||
# ngram is not supported now
|
||||
ngram = None
|
||||
scorers["ngram"] = ngram
|
||||
|
||||
weights = dict(
|
||||
decoder=1.0 - kwargs.get("decoding_ctc_weight", 0.0),
|
||||
ctc=kwargs.get("decoding_ctc_weight", 0.0),
|
||||
lm=kwargs.get("lm_weight", 0.0),
|
||||
ngram=kwargs.get("ngram_weight", 0.0),
|
||||
length_bonus=kwargs.get("penalty", 0.0),
|
||||
)
|
||||
|
||||
beam_search = BeamSearchScamaStreaming(
|
||||
beam_size=kwargs.get("beam_size", 2),
|
||||
weights=weights,
|
||||
scorers=scorers,
|
||||
sos=self.sos,
|
||||
eos=self.eos,
|
||||
vocab_size=len(token_list),
|
||||
token_list=token_list,
|
||||
pre_beam_score_key=None if self.ctc_weight == 1.0 else "full",
|
||||
)
|
||||
# beam_search.to(device=kwargs.get("device", "cpu"), dtype=getattr(torch, kwargs.get("dtype", "float32"))).eval()
|
||||
# for scorer in scorers.values():
|
||||
# if isinstance(scorer, torch.nn.Module):
|
||||
# scorer.to(device=kwargs.get("device", "cpu"), dtype=getattr(torch, kwargs.get("dtype", "float32"))).eval()
|
||||
self.beam_search = beam_search
|
||||
|
||||
def generate_chunk(
|
||||
self,
|
||||
speech,
|
||||
speech_lengths=None,
|
||||
key: list = None,
|
||||
tokenizer=None,
|
||||
frontend=None,
|
||||
**kwargs,
|
||||
):
|
||||
cache = kwargs.get("cache", {})
|
||||
speech = speech.to(device=kwargs["device"])
|
||||
speech_lengths = speech_lengths.to(device=kwargs["device"])
|
||||
|
||||
# Encoder
|
||||
encoder_out, encoder_out_lens = self.encode_chunk(
|
||||
speech, speech_lengths, cache=cache, is_final=kwargs.get("is_final", False)
|
||||
)
|
||||
if isinstance(encoder_out, tuple):
|
||||
encoder_out = encoder_out[0]
|
||||
if "running_hyps" not in cache:
|
||||
running_hyps = self.beam_search.init_hyp(encoder_out)
|
||||
cache["running_hyps"] = running_hyps
|
||||
|
||||
# predictor
|
||||
predictor_outs = self.calc_predictor_chunk(
|
||||
encoder_out,
|
||||
encoder_out_lens,
|
||||
cache=cache,
|
||||
is_final=kwargs.get("is_final", False),
|
||||
)
|
||||
pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = (
|
||||
predictor_outs[0],
|
||||
predictor_outs[1],
|
||||
predictor_outs[2],
|
||||
predictor_outs[3],
|
||||
)
|
||||
pre_token_length = pre_token_length.round().long()
|
||||
|
||||
if torch.max(pre_token_length) < 1:
|
||||
return []
|
||||
maxlen = minlen = pre_token_length
|
||||
if kwargs.get("is_final", False):
|
||||
maxlen += kwargs.get("token_num_relax", 5)
|
||||
minlen = max(0, minlen - kwargs.get("token_num_relax", 5))
|
||||
# c. Passed the encoder result and the beam search
|
||||
nbest_hyps = self.beam_search(
|
||||
x=encoder_out[0],
|
||||
scama_mask=None,
|
||||
pre_acoustic_embeds=pre_acoustic_embeds,
|
||||
maxlen=int(maxlen),
|
||||
minlen=int(minlen),
|
||||
cache=cache,
|
||||
)
|
||||
|
||||
cache["running_hyps"] = nbest_hyps
|
||||
nbest_hyps = nbest_hyps[: self.nbest]
|
||||
|
||||
results = []
|
||||
for hyp in nbest_hyps:
|
||||
# assert isinstance(hyp, (Hypothesis)), type(hyp)
|
||||
|
||||
# remove sos/eos and get results
|
||||
last_pos = -1
|
||||
if isinstance(hyp.yseq, list):
|
||||
token_int = hyp.yseq[1:last_pos]
|
||||
else:
|
||||
token_int = hyp.yseq[1:last_pos].tolist()
|
||||
|
||||
# remove blank symbol id, which is assumed to be 0
|
||||
token_int = list(
|
||||
filter(
|
||||
lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int
|
||||
)
|
||||
)
|
||||
|
||||
# Change integer-ids to tokens
|
||||
token = tokenizer.ids2tokens(token_int)
|
||||
# text = tokenizer.tokens2text(token)
|
||||
|
||||
result_i = token
|
||||
|
||||
results.extend(result_i)
|
||||
|
||||
return results
|
||||
|
||||
def init_cache(self, cache: dict = {}, **kwargs):
|
||||
device = kwargs.get("device", "cuda")
|
||||
|
||||
chunk_size = kwargs.get("chunk_size", [0, 10, 5])
|
||||
encoder_chunk_look_back = kwargs.get("encoder_chunk_look_back", 0)
|
||||
decoder_chunk_look_back = kwargs.get("decoder_chunk_look_back", 0)
|
||||
batch_size = 1
|
||||
|
||||
enc_output_size = kwargs["encoder_conf"]["output_size"]
|
||||
feats_dims = kwargs["frontend_conf"]["n_mels"] * kwargs["frontend_conf"]["lfr_m"]
|
||||
|
||||
cache_encoder = {
|
||||
"start_idx": 0,
|
||||
"cif_hidden": torch.zeros((batch_size, 1, enc_output_size)).to(device=device),
|
||||
"cif_alphas": torch.zeros((batch_size, 1)).to(device=device),
|
||||
"chunk_size": chunk_size,
|
||||
"encoder_chunk_look_back": encoder_chunk_look_back,
|
||||
"last_chunk": False,
|
||||
"opt": None,
|
||||
"feats": torch.zeros((batch_size, chunk_size[0] + chunk_size[2], feats_dims)).to(
|
||||
device=device
|
||||
),
|
||||
"tail_chunk": False,
|
||||
}
|
||||
cache["encoder"] = cache_encoder
|
||||
|
||||
cache_decoder = {
|
||||
"decode_fsmn": None,
|
||||
"decoder_chunk_look_back": decoder_chunk_look_back,
|
||||
"opt": None,
|
||||
"chunk_size": chunk_size,
|
||||
}
|
||||
cache["decoder"] = cache_decoder
|
||||
cache["frontend"] = {}
|
||||
|
||||
cache["prev_samples"] = torch.empty(0).to(device=device)
|
||||
|
||||
return cache
|
||||
|
||||
def inference(
|
||||
self,
|
||||
data_in,
|
||||
data_lengths=None,
|
||||
key: list = None,
|
||||
tokenizer=None,
|
||||
frontend=None,
|
||||
cache: dict = {},
|
||||
**kwargs,
|
||||
):
|
||||
|
||||
# init beamsearch
|
||||
is_use_ctc = kwargs.get("decoding_ctc_weight", 0.0) > 0.00001 and self.ctc != None
|
||||
is_use_lm = (
|
||||
kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None
|
||||
)
|
||||
|
||||
if self.beam_search is None:
|
||||
|
||||
logging.info("enable beam_search")
|
||||
self.init_beam_search(**kwargs)
|
||||
self.nbest = kwargs.get("nbest", 1)
|
||||
|
||||
if len(cache) == 0:
|
||||
self.init_cache(cache, **kwargs)
|
||||
|
||||
meta_data = {}
|
||||
chunk_size = kwargs.get("chunk_size", [0, 10, 5])
|
||||
chunk_stride_samples = int(chunk_size[1] * 960) # 600ms
|
||||
|
||||
time1 = time.perf_counter()
|
||||
cfg = {"is_final": kwargs.get("is_final", False)}
|
||||
audio_sample_list = load_audio_text_image_video(
|
||||
data_in,
|
||||
fs=frontend.fs,
|
||||
audio_fs=kwargs.get("fs", 16000),
|
||||
data_type=kwargs.get("data_type", "sound"),
|
||||
tokenizer=tokenizer,
|
||||
cache=cfg,
|
||||
)
|
||||
_is_final = cfg["is_final"] # if data_in is a file or url, set is_final=True
|
||||
|
||||
time2 = time.perf_counter()
|
||||
meta_data["load_data"] = f"{time2 - time1:0.3f}"
|
||||
assert len(audio_sample_list) == 1, "batch_size must be set 1"
|
||||
|
||||
audio_sample = torch.cat((cache["prev_samples"], audio_sample_list[0]))
|
||||
|
||||
n = int(len(audio_sample) // chunk_stride_samples + int(_is_final))
|
||||
m = int(len(audio_sample) % chunk_stride_samples * (1 - int(_is_final)))
|
||||
tokens = []
|
||||
for i in range(n):
|
||||
kwargs["is_final"] = _is_final and i == n - 1
|
||||
audio_sample_i = audio_sample[i * chunk_stride_samples : (i + 1) * chunk_stride_samples]
|
||||
|
||||
# extract fbank feats
|
||||
speech, speech_lengths = extract_fbank(
|
||||
[audio_sample_i],
|
||||
data_type=kwargs.get("data_type", "sound"),
|
||||
frontend=frontend,
|
||||
cache=cache["frontend"],
|
||||
is_final=kwargs["is_final"],
|
||||
)
|
||||
time3 = time.perf_counter()
|
||||
meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
|
||||
meta_data["batch_data_time"] = (
|
||||
speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000
|
||||
)
|
||||
|
||||
tokens_i = self.generate_chunk(
|
||||
speech,
|
||||
speech_lengths,
|
||||
key=key,
|
||||
tokenizer=tokenizer,
|
||||
cache=cache,
|
||||
frontend=frontend,
|
||||
**kwargs,
|
||||
)
|
||||
tokens.extend(tokens_i)
|
||||
|
||||
text_postprocessed, _ = postprocess_utils.sentence_postprocess(tokens)
|
||||
|
||||
result_i = {"key": key[0], "text": text_postprocessed}
|
||||
result = [result_i]
|
||||
|
||||
cache["prev_samples"] = audio_sample[:-m]
|
||||
if _is_final:
|
||||
self.init_cache(cache, **kwargs)
|
||||
|
||||
if kwargs.get("output_dir"):
|
||||
writer = DatadirWriter(kwargs.get("output_dir"))
|
||||
ibest_writer = writer[f"{1}best_recog"]
|
||||
ibest_writer["token"][key[0]] = " ".join(tokens)
|
||||
ibest_writer["text"][key[0]] = text_postprocessed
|
||||
|
||||
return result, meta_data
|
||||
@@ -0,0 +1,127 @@
|
||||
# This is an example that demonstrates how to configure a model file.
|
||||
# You can modify the configuration according to your own requirements.
|
||||
|
||||
# to print the register_table:
|
||||
# from funasr.register import tables
|
||||
# tables.print()
|
||||
|
||||
# network architecture
|
||||
model: SCAMA
|
||||
model_conf:
|
||||
ctc_weight: 0.0
|
||||
lsm_weight: 0.1
|
||||
length_normalized_loss: true
|
||||
|
||||
# encoder
|
||||
encoder: SANMEncoderChunkOpt
|
||||
encoder_conf:
|
||||
output_size: 512
|
||||
attention_heads: 4
|
||||
linear_units: 2048
|
||||
num_blocks: 50
|
||||
dropout_rate: 0.1
|
||||
positional_dropout_rate: 0.1
|
||||
attention_dropout_rate: 0.1
|
||||
input_layer: pe
|
||||
pos_enc_class: SinusoidalPositionEncoder
|
||||
normalize_before: true
|
||||
kernel_size: 11
|
||||
sanm_shfit: 0
|
||||
selfattention_layer_type: sanm
|
||||
|
||||
# decoder
|
||||
decoder: FsmnDecoderSCAMAOpt
|
||||
decoder_conf:
|
||||
attention_heads: 4
|
||||
linear_units: 2048
|
||||
num_blocks: 16
|
||||
dropout_rate: 0.1
|
||||
positional_dropout_rate: 0.1
|
||||
self_attention_dropout_rate: 0.1
|
||||
src_attention_dropout_rate: 0.1
|
||||
att_layer_num: 16
|
||||
kernel_size: 11
|
||||
sanm_shfit: 0
|
||||
|
||||
predictor: CifPredictorV2
|
||||
predictor_conf:
|
||||
idim: 512
|
||||
threshold: 1.0
|
||||
l_order: 1
|
||||
r_order: 1
|
||||
tail_threshold: 0.45
|
||||
|
||||
# frontend related
|
||||
frontend: WavFrontend
|
||||
frontend_conf:
|
||||
fs: 16000
|
||||
window: hamming
|
||||
n_mels: 80
|
||||
frame_length: 25
|
||||
frame_shift: 10
|
||||
lfr_m: 7
|
||||
lfr_n: 6
|
||||
|
||||
specaug: SpecAugLFR
|
||||
specaug_conf:
|
||||
apply_time_warp: false
|
||||
time_warp_window: 5
|
||||
time_warp_mode: bicubic
|
||||
apply_freq_mask: true
|
||||
freq_mask_width_range:
|
||||
- 0
|
||||
- 30
|
||||
lfr_rate: 6
|
||||
num_freq_mask: 1
|
||||
apply_time_mask: true
|
||||
time_mask_width_range:
|
||||
- 0
|
||||
- 12
|
||||
num_time_mask: 1
|
||||
|
||||
train_conf:
|
||||
accum_grad: 1
|
||||
grad_clip: 5
|
||||
max_epoch: 150
|
||||
val_scheduler_criterion:
|
||||
- valid
|
||||
- acc
|
||||
best_model_criterion:
|
||||
- - valid
|
||||
- acc
|
||||
- max
|
||||
keep_nbest_models: 10
|
||||
avg_nbest_model: 10
|
||||
log_interval: 50
|
||||
|
||||
optim: adam
|
||||
optim_conf:
|
||||
lr: 0.0005
|
||||
scheduler: warmuplr
|
||||
scheduler_conf:
|
||||
warmup_steps: 30000
|
||||
|
||||
dataset: AudioDataset
|
||||
dataset_conf:
|
||||
index_ds: IndexDSJsonl
|
||||
batch_sampler: BatchSampler
|
||||
batch_type: example # example or length
|
||||
batch_size: 1 # if batch_type is example, batch_size is the numbers of samples; if length, batch_size is source_token_len+target_token_len;
|
||||
max_token_length: 2048 # filter samples if source_token_len+target_token_len > max_token_length,
|
||||
buffer_size: 500
|
||||
shuffle: True
|
||||
num_workers: 0
|
||||
|
||||
tokenizer: CharTokenizer
|
||||
tokenizer_conf:
|
||||
unk_symbol: <unk>
|
||||
split_with_space: true
|
||||
|
||||
|
||||
ctc_conf:
|
||||
dropout_rate: 0.0
|
||||
ctc_type: builtin
|
||||
reduce: true
|
||||
ignore_nan_grad: true
|
||||
|
||||
normalize: null
|
||||
@@ -0,0 +1,89 @@
|
||||
import os
|
||||
import yaml
|
||||
import torch
|
||||
import numpy as np
|
||||
from torch.nn import functional as F
|
||||
|
||||
|
||||
def sequence_mask(lengths, maxlen=None, dtype=torch.float32, device=None):
|
||||
if maxlen is None:
|
||||
maxlen = lengths.max()
|
||||
row_vector = torch.arange(0, maxlen, 1).to(lengths.device)
|
||||
matrix = torch.unsqueeze(lengths, dim=-1)
|
||||
mask = row_vector < matrix
|
||||
mask = mask.detach()
|
||||
|
||||
return mask.type(dtype).to(device) if device is not None else mask.type(dtype)
|
||||
|
||||
|
||||
def apply_cmvn(inputs, mvn):
|
||||
device = inputs.device
|
||||
dtype = inputs.dtype
|
||||
frame, dim = inputs.shape
|
||||
meams = np.tile(mvn[0:1, :dim], (frame, 1))
|
||||
vars = np.tile(mvn[1:2, :dim], (frame, 1))
|
||||
inputs -= torch.from_numpy(meams).type(dtype).to(device)
|
||||
inputs *= torch.from_numpy(vars).type(dtype).to(device)
|
||||
|
||||
return inputs.type(torch.float32)
|
||||
|
||||
|
||||
def drop_and_add(
|
||||
inputs: torch.Tensor,
|
||||
outputs: torch.Tensor,
|
||||
training: bool,
|
||||
dropout_rate: float = 0.1,
|
||||
stoch_layer_coeff: float = 1.0,
|
||||
):
|
||||
|
||||
outputs = F.dropout(outputs, p=dropout_rate, training=training, inplace=True)
|
||||
outputs *= stoch_layer_coeff
|
||||
|
||||
input_dim = inputs.size(-1)
|
||||
output_dim = outputs.size(-1)
|
||||
|
||||
if input_dim == output_dim:
|
||||
outputs += inputs
|
||||
return outputs
|
||||
|
||||
|
||||
def proc_tf_vocab(vocab_path):
|
||||
with open(vocab_path, encoding="utf-8") as f:
|
||||
token_list = [line.rstrip() for line in f]
|
||||
if "<unk>" not in token_list:
|
||||
token_list.append("<unk>")
|
||||
return token_list
|
||||
|
||||
|
||||
def gen_config_for_tfmodel(config_path, vocab_path, output_dir):
|
||||
token_list = proc_tf_vocab(vocab_path)
|
||||
with open(config_path, encoding="utf-8") as f:
|
||||
config = yaml.safe_load(f)
|
||||
|
||||
config["token_list"] = token_list
|
||||
|
||||
if not os.path.exists(output_dir):
|
||||
os.makedirs(output_dir)
|
||||
|
||||
with open(os.path.join(output_dir, "config.yaml"), "w", encoding="utf-8") as f:
|
||||
yaml_no_alias_safe_dump(config, f, indent=4, sort_keys=False)
|
||||
|
||||
|
||||
class NoAliasSafeDumper(yaml.SafeDumper):
|
||||
# Disable anchor/alias in yaml because looks ugly
|
||||
def ignore_aliases(self, data):
|
||||
return True
|
||||
|
||||
|
||||
def yaml_no_alias_safe_dump(data, stream=None, **kwargs):
|
||||
"""Safe-dump in yaml with no anchor/alias"""
|
||||
return yaml.dump(data, stream, allow_unicode=True, Dumper=NoAliasSafeDumper, **kwargs)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
config_path = sys.argv[1]
|
||||
vocab_path = sys.argv[2]
|
||||
output_dir = sys.argv[3]
|
||||
gen_config_for_tfmodel(config_path, vocab_path, output_dir)
|
||||
Reference in new issue
Block a user