From d46a542fae26009eee16204a81903862cb4dba73 Mon Sep 17 00:00:00 2001
From: aky15 <ankeyu.aky@11.17.44.249>
Date: 星期一, 10 四月 2023 16:02:41 +0800
Subject: [PATCH] Merge branch 'main' into dev_aky
---
funasr/train/abs_model.py | 192 ++++++++++++++++++++++++++++++++++++++++++++++++
1 files changed, 192 insertions(+), 0 deletions(-)
diff --git a/funasr/train/abs_model.py b/funasr/train/abs_model.py
new file mode 100644
index 0000000..1c7ff3d
--- /dev/null
+++ b/funasr/train/abs_model.py
@@ -0,0 +1,192 @@
+from abc import ABC
+from abc import abstractmethod
+
+
+from typing import Dict
+from typing import Optional
+from typing import Tuple
+
+import torch
+import torch.nn.functional as F
+from typeguard import check_argument_types
+
+from funasr.modules.nets_utils import make_pad_mask
+from funasr.torch_utils.device_funcs import force_gatherable
+from funasr.train.abs_espnet_model import AbsESPnetModel
+
+from funasr.modules.scorers.scorer_interface import BatchScorerInterface
+
+
+class AbsPunctuation(torch.nn.Module, BatchScorerInterface, ABC):
+ """The abstract class
+
+ To share the loss calculation way among different models,
+ We uses delegate pattern here:
+ The instance of this class should be passed to "LanguageModel"
+
+ This "model" is one of mediator objects for "Task" class.
+
+ """
+
+ @abstractmethod
+ def forward(self, input: torch.Tensor, hidden: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ raise NotImplementedError
+
+ @abstractmethod
+ def with_vad(self) -> bool:
+ raise NotImplementedError
+
+
+class PunctuationModel(AbsESPnetModel):
+
+ def __init__(self, punc_model: AbsPunctuation, vocab_size: int, ignore_id: int = 0, punc_weight: list = None):
+ assert check_argument_types()
+ super().__init__()
+ self.punc_model = punc_model
+ self.punc_weight = torch.Tensor(punc_weight)
+ self.sos = 1
+ self.eos = 2
+
+ # ignore_id may be assumed as 0, shared with CTC-blank symbol for ASR.
+ self.ignore_id = ignore_id
+ # if self.punc_model.with_vad():
+ # print("This is a vad puncuation model.")
+
+ def nll(
+ self,
+ text: torch.Tensor,
+ punc: torch.Tensor,
+ text_lengths: torch.Tensor,
+ punc_lengths: torch.Tensor,
+ max_length: Optional[int] = None,
+ vad_indexes: Optional[torch.Tensor] = None,
+ vad_indexes_lengths: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Compute negative log likelihood(nll)
+
+ Normally, this function is called in batchify_nll.
+ Args:
+ text: (Batch, Length)
+ punc: (Batch, Length)
+ text_lengths: (Batch,)
+ max_lengths: int
+ """
+ batch_size = text.size(0)
+ # For data parallel
+ if max_length is None:
+ text = text[:, :text_lengths.max()]
+ punc = punc[:, :text_lengths.max()]
+ else:
+ text = text[:, :max_length]
+ punc = punc[:, :max_length]
+
+ if self.punc_model.with_vad():
+ # Should be VadRealtimeTransformer
+ assert vad_indexes is not None
+ y, _ = self.punc_model(text, text_lengths, vad_indexes)
+ else:
+ # Should be TargetDelayTransformer,
+ y, _ = self.punc_model(text, text_lengths)
+
+ # Calc negative log likelihood
+ # nll: (BxL,)
+ if self.training == False:
+ _, indices = y.view(-1, y.shape[-1]).topk(1, dim=1)
+ from sklearn.metrics import f1_score
+ f1_score = f1_score(punc.view(-1).detach().cpu().numpy(),
+ indices.squeeze(-1).detach().cpu().numpy(),
+ average='micro')
+ nll = torch.Tensor([f1_score]).repeat(text_lengths.sum())
+ return nll, text_lengths
+ else:
+ self.punc_weight = self.punc_weight.to(punc.device)
+ nll = F.cross_entropy(y.view(-1, y.shape[-1]), punc.view(-1), self.punc_weight, reduction="none",
+ ignore_index=self.ignore_id)
+ # nll: (BxL,) -> (BxL,)
+ if max_length is None:
+ nll.masked_fill_(make_pad_mask(text_lengths).to(nll.device).view(-1), 0.0)
+ else:
+ nll.masked_fill_(
+ make_pad_mask(text_lengths, maxlen=max_length + 1).to(nll.device).view(-1),
+ 0.0,
+ )
+ # nll: (BxL,) -> (B, L)
+ nll = nll.view(batch_size, -1)
+ return nll, text_lengths
+
+ def batchify_nll(self,
+ text: torch.Tensor,
+ punc: torch.Tensor,
+ text_lengths: torch.Tensor,
+ punc_lengths: torch.Tensor,
+ batch_size: int = 100) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Compute negative log likelihood(nll) from transformer language model
+
+ To avoid OOM, this fuction seperate the input into batches.
+ Then call nll for each batch and combine and return results.
+ Args:
+ text: (Batch, Length)
+ punc: (Batch, Length)
+ text_lengths: (Batch,)
+ batch_size: int, samples each batch contain when computing nll,
+ you may change this to avoid OOM or increase
+
+ """
+ total_num = text.size(0)
+ if total_num <= batch_size:
+ nll, x_lengths = self.nll(text, punc, text_lengths)
+ else:
+ nlls = []
+ x_lengths = []
+ max_length = text_lengths.max()
+
+ start_idx = 0
+ while True:
+ end_idx = min(start_idx + batch_size, total_num)
+ batch_text = text[start_idx:end_idx, :]
+ batch_punc = punc[start_idx:end_idx, :]
+ batch_text_lengths = text_lengths[start_idx:end_idx]
+ # batch_nll: [B * T]
+ batch_nll, batch_x_lengths = self.nll(batch_text, batch_punc, batch_text_lengths, max_length=max_length)
+ nlls.append(batch_nll)
+ x_lengths.append(batch_x_lengths)
+ start_idx = end_idx
+ if start_idx == total_num:
+ break
+ nll = torch.cat(nlls)
+ x_lengths = torch.cat(x_lengths)
+ assert nll.size(0) == total_num
+ assert x_lengths.size(0) == total_num
+ return nll, x_lengths
+
+ def forward(
+ self,
+ text: torch.Tensor,
+ punc: torch.Tensor,
+ text_lengths: torch.Tensor,
+ punc_lengths: torch.Tensor,
+ vad_indexes: Optional[torch.Tensor] = None,
+ vad_indexes_lengths: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]:
+ nll, y_lengths = self.nll(text, punc, text_lengths, punc_lengths, vad_indexes=vad_indexes)
+ ntokens = y_lengths.sum()
+ loss = nll.sum() / ntokens
+ stats = dict(loss=loss.detach())
+
+ # force_gatherable: to-device and to-tensor if scalar for DataParallel
+ loss, stats, weight = force_gatherable((loss, stats, ntokens), loss.device)
+ return loss, stats, weight
+
+ def collect_feats(self, text: torch.Tensor, punc: torch.Tensor,
+ text_lengths: torch.Tensor) -> Dict[str, torch.Tensor]:
+ return {}
+
+ def inference(self,
+ text: torch.Tensor,
+ text_lengths: torch.Tensor,
+ vad_indexes: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, None]:
+ if self.punc_model.with_vad():
+ assert vad_indexes is not None
+ return self.punc_model(text, text_lengths, vad_indexes)
+ else:
+ return self.punc_model(text, text_lengths)
--
Gitblit v1.9.1