Initial commit: FunASR Speech Recognition Toolkit

Add complete FunASR codebase including models, runtime, and documentation.

funasr/models/paraformer_v2_community/model.py

@@ -0,0 +1,603 @@
+#!/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 copy
+import torch
+import logging
+from torch.cuda.amp import autocast
+from typing import Union, Dict, List, Tuple, Optional
+
+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.train_utils.device_funcs import to_device
+from funasr.utils.datadir_writer import DatadirWriter
+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
+from funasr.utils.timestamp_tools import ts_prediction_lfr6_standard
+from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
+from torch.nn.utils.rnn import pad_sequence
+import torchaudio
+
+@tables.register("model_classes", "Paraformer_v2_community")
+class Paraformer(torch.nn.Module):
+    """
+    Author: Speech Lab of DAMO Academy, Alibaba Group
+    Paraformer: Fast and Accurate Parallel Transformer for Non-autoregressive End-to-End Speech Recognition
+    https://arxiv.org/abs/2206.08317
+    """
+
+    def __init__(
+        self,
+        specaug: Optional[str] = None,
+        specaug_conf: Optional[Dict] = None,
+        normalize: str = None,
+        normalize_conf: Optional[Dict] = None,
+        encoder: str = None,
+        encoder_conf: Optional[Dict] = None,
+        decoder: str = None,
+        decoder_conf: Optional[Dict] = None,
+        ctc: str = None,
+        ctc_conf: Optional[Dict] = None,
+        ctc_weight: float = 0.5,
+        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,
+        # report_cer: bool = True,
+        # report_wer: bool = True,
+        # sym_space: str = "<space>",
+        # sym_blank: str = "<blank>",
+        # extract_feats_in_collect_stats: bool = True,
+        # predictor=None,
+        share_embedding: bool = False,
+        # preencoder: Optional[AbsPreEncoder] = None,
+        # postencoder: Optional[AbsPostEncoder] = None,
+        use_1st_decoder_loss: bool = False,
+        **kwargs,
+    ):
+
+        """Initialize Paraformer.
+        
+            Args:
+                specaug: TODO.
+                specaug_conf: Configuration dict for specaug.
+                normalize: TODO.
+                normalize_conf: Configuration dict for normalize.
+                encoder: TODO.
+                encoder_conf: Configuration dict for encoder.
+                decoder: TODO.
+                decoder_conf: Configuration dict for decoder.
+                ctc: TODO.
+                ctc_conf: Configuration dict for ctc.
+                ctc_weight: TODO.
+                input_size: Size/dimension parameter.
+                vocab_size: Size/dimension parameter.
+                ignore_id: TODO.
+                blank_id: TODO.
+                sos: TODO.
+                eos: TODO.
+                lsm_weight: TODO.
+                length_normalized_loss: TODO.
+                share_embedding: TODO.
+                use_1st_decoder_loss: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        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()
+
+        if decoder is not None:
+            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)
+
+        # 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.token_list = token_list.copy()
+        #
+        # self.frontend = frontend
+        self.specaug = specaug
+        self.normalize = normalize
+        # self.preencoder = preencoder
+        # self.postencoder = postencoder
+        self.encoder = encoder
+        #
+        # if not hasattr(self.encoder, "interctc_use_conditioning"):
+        #     self.encoder.interctc_use_conditioning = False
+        # if self.encoder.interctc_use_conditioning:
+        #     self.encoder.conditioning_layer = torch.nn.Linear(
+        #         vocab_size, self.encoder.output_size()
+        #     )
+        #
+        # self.error_calculator = None
+        #
+        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 report_cer or report_wer:
+        #     self.error_calculator = ErrorCalculator(
+        #         token_list, sym_space, sym_blank, report_cer, report_wer
+        #     )
+        #
+        if ctc_weight == 0.0:
+            self.ctc = None
+        else:
+            self.ctc = ctc
+        #
+        # self.extract_feats_in_collect_stats = extract_feats_in_collect_stats
+
+        self.share_embedding = share_embedding
+        if self.share_embedding:
+            self.decoder.embed = None
+
+        self.use_1st_decoder_loss = use_1st_decoder_loss
+        self.length_normalized_loss = length_normalized_loss
+        self.beam_search = None
+        self.error_calculator = None
+
+    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,)
+        """
+        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
+        encoder_out, encoder_out_lens = self.encode(speech, speech_lengths)
+
+        loss_ctc, cer_ctc = None, None
+        loss_pre = None
+        stats = dict()
+
+        # decoder: CTC branch
+        if self.ctc_weight != 0.0:
+            loss_ctc, cer_ctc = self._calc_ctc_loss(
+                encoder_out, encoder_out_lens, 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 = self._calc_att_loss(
+            encoder_out, encoder_out_lens, text, text_lengths
+        )
+
+        # 3. CTC-Att loss definition
+        if self.ctc_weight == 0.0:
+            loss = loss_att
+        else:
+            loss = (
+                self.ctc_weight * loss_ctc
+                + (1 - self.ctc_weight) * loss_att
+            )
+
+        # 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"] = torch.clone(loss.detach())
+        stats["batch_size"] = batch_size
+
+        # force_gatherable: to-device and to-tensor if scalar for DataParallel
+        if self.length_normalized_loss:
+            batch_size = (text_lengths).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 _calc_att_loss(
+        self,
+        encoder_out: torch.Tensor,
+        encoder_out_lens: torch.Tensor,
+        ys_pad: torch.Tensor,
+        ys_pad_lens: torch.Tensor,
+    ):
+
+        # 0. sampler
+        """Internal: calc att loss.
+        
+            Args:
+                encoder_out: Encoder output tensor.
+                encoder_out_lens: Encoder output lengths.
+                ys_pad: TODO.
+                ys_pad_lens: Lengths of ys_pad.
+            """
+        decoder_out_1st = None
+
+        batch_size = encoder_out.size(0)
+        ctc_probs_all = self.ctc.softmax(encoder_out)
+        compressed_ctc_list = []
+        for b in range(batch_size):
+            ctc_prob_b = ctc_probs_all[b, :encoder_out_lens[b]]
+            text_b = ys_pad[b, :ys_pad_lens[b]]
+            with torch.no_grad():
+                ctc_log_prob_b = ctc_prob_b.log()
+                align_path = self.force_align(ctc_log_prob_b.cpu(), text_b.cpu(), blank_id=self.blank_id)
+                align_path = align_path.to(encoder_out.device)
+                target_idx_path = self.map_alignment_to_target_index(align_path, self.blank_id)
+            ctc_comp = self.average_repeats_training(ctc_prob_b, target_idx_path, ys_pad_lens[b])
+            compressed_ctc_list.append(ctc_comp)
+
+        # 4. Pad Batch to [B, U_max, V]
+        padded_ctc_input = pad_sequence(compressed_ctc_list, batch_first=True).to(encoder_out.device)
+
+
+
+        # 1. Forward decoder
+        decoder_outs = self.decoder(encoder_out, encoder_out_lens, padded_ctc_input, ys_pad_lens)
+        decoder_out, _ = decoder_outs[0], decoder_outs[1]
+
+        if decoder_out_1st is None:
+            decoder_out_1st = decoder_out
+        # 2. Compute attention loss
+        loss_att = self.criterion_att(decoder_out, ys_pad)
+        acc_att = th_accuracy(
+            decoder_out_1st.view(-1, self.vocab_size),
+            ys_pad,
+            ignore_label=self.ignore_id,
+        )
+
+        # 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_1st.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
+
+
+    def _calc_ctc_loss(
+        self,
+        encoder_out: torch.Tensor,
+        encoder_out_lens: torch.Tensor,
+        ys_pad: torch.Tensor,
+        ys_pad_lens: torch.Tensor,
+    ):
+        # Calc CTC loss
+        """Internal: calc ctc loss.
+        
+            Args:
+                encoder_out: Encoder output tensor.
+                encoder_out_lens: Encoder output lengths.
+                ys_pad: TODO.
+                ys_pad_lens: Lengths of ys_pad.
+            """
+        loss_ctc = self.ctc(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens)
+
+        # Calc CER using CTC
+        cer_ctc = None
+        if not self.training and self.error_calculator is not None:
+            ys_hat = self.ctc.argmax(encoder_out).data
+            cer_ctc = self.error_calculator(ys_hat.cpu(), ys_pad.cpu(), is_ctc=True)
+        return loss_ctc, cer_ctc
+    
+    def map_alignment_to_target_index(self, align_path, blank_id):
+        """
+        Robustly map CTC alignment path (Token IDs) to Target Indices.
+        
+        Logic:
+            Detect boundaries where a new token segment begins.
+            A segment starts if the current frame is a Token AND it is different from the previous frame 
+            (considering CTC topology where repeats are separated by blanks or are distinct tokens).
+        
+        Example:
+            Text: [A, B]
+            Align Path: [A, A, _, B, B]
+            Output:     [0, 0, -1, 1, 1]
+        """
+        # 1. Identify where the path is NOT blank
+        is_token = align_path != blank_id
+        
+        # 2. Identify transitions
+        prev_path = torch.roll(align_path, 1)
+        # Handle the very first frame: if it's a token, it must be the start of segment 0.
+        prev_path[0] = blank_id # force mismatch for the first element
+        
+        # A new segment starts if: It's a token AND (it differs from prev OR prev was blank)
+        # Note: If align_path[i] == align_path[i-1] (and not blank), it's the same segment.
+        new_segment_start = is_token & (align_path != prev_path)
+        
+        # 3. Cumulative sum to assign indices (1..U)
+        segment_ids = torch.cumsum(new_segment_start.long(), dim=0) - 1
+        
+        # 4. Mask out blank positions with -1
+        target_idx_path = torch.where(is_token, segment_ids, -1)
+        
+        return target_idx_path
+    def force_align(self, ctc_probs: torch.Tensor, y: torch.Tensor, blank_id=0) -> list:
+        """ctc forced alignment.
+
+        Args:
+            torch.Tensor ctc_probs: hidden state sequence, 2d tensor (T, D)
+            torch.Tensor y: id sequence tensor 1d tensor (L)
+            int blank_id: blank symbol index
+        Returns:
+            torch.Tensor: alignment result
+        """
+        ctc_probs = ctc_probs[None].cpu()
+        y = y[None].cpu()
+        alignments, _ = torchaudio.functional.forced_align(ctc_probs, y, blank=blank_id)
+        return alignments[0]
+
+    def average_repeats_training(self, ctc_probs, target_idx_path, target_len):
+        """
+        Aggregates frames belonging to the same target index using scatter_add.
+        
+        Args:
+            ctc_probs: [T, V]
+            target_idx_path: [T], values in [-1, 0, ... U-1]
+            target_len: U
+        Returns:
+            compressed: [U, V]
+        """
+        U = target_len
+        V = ctc_probs.size(1)
+        
+        compressed = torch.zeros((U, V), device=ctc_probs.device, dtype=ctc_probs.dtype)
+        counts = torch.zeros((U, 1), device=ctc_probs.device, dtype=ctc_probs.dtype)
+        
+        # Filter valid frames (non-blank)
+        mask = target_idx_path != -1
+        valid_indices = target_idx_path[mask] # [T_valid]
+        valid_probs = ctc_probs[mask]         # [T_valid, V]
+        
+        if valid_indices.numel() == 0:
+            return compressed
+            
+        # Scatter Add Probs
+        index_expanded = valid_indices.unsqueeze(1).repeat(1, V)
+        compressed.scatter_add_(0, index_expanded, valid_probs)
+        
+        # Scatter Add Counts
+        ones = torch.ones((valid_indices.size(0), 1), device=ctc_probs.device)
+        counts.scatter_add_(0, valid_indices.unsqueeze(1), ones)
+        
+        # Average
+        compressed = compressed / (counts + 1e-9)
+        return compressed
+    
+    def average_repeats_inference(self, ctc_probs, greedy_path):
+        """
+        Returns:
+            merged_probs: [U', V]
+            timestamps: List[Tuple[int, int]] -> [(start_frame, end_frame), ...]
+        """
+        if greedy_path.numel() == 0:
+            return torch.zeros((0, ctc_probs.size(1)), device=ctc_probs.device)
+
+        # Find consecutive segments in the greedy path
+        unique_tokens, counts = torch.unique_consecutive(greedy_path, return_counts=True)
+        
+        # Compute start and end indices for each segment
+        end_indices = torch.cumsum(counts, dim=0)
+        start_indices = torch.cat([torch.tensor([0], device=counts.device), end_indices[:-1]])
+        
+        merged_probs = []
+    
+        for i, token in enumerate(unique_tokens):
+            if token != self.blank_id:
+                start = start_indices[i].item()
+                end = end_indices[i].item()
+                
+                # Extract and average probabilities for the decoder
+                avg_prob = ctc_probs[start:end].mean(dim=0)
+                merged_probs.append(avg_prob)
+
+        
+        if not merged_probs:
+            return torch.zeros((0, ctc_probs.size(1)), device=ctc_probs.device)
+            
+        return torch.stack(merged_probs)
+
+    def inference(
+        self,
+        data_in,
+        data_lengths=None,
+        key: list = None,
+        tokenizer=None,
+        frontend=None,
+        **kwargs,
+    ):
+
+        """Run inference on input data.
+        
+            Args:
+                data_in: Input data (audio samples, file paths, or text).
+                data_lengths: Lengths of each input sample in the batch.
+                key: Sample identifiers.
+                tokenizer: Tokenizer instance for text encoding/decoding.
+                frontend: Audio frontend for feature extraction.
+                **kwargs: Additional keyword arguments.
+            """
+        meta_data = {}
+        if (
+            isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank"
+        ):  # fbank
+            speech, speech_lengths = data_in, data_lengths
+            if len(speech.shape) < 3:
+                speech = speech[None, :, :]
+            if speech_lengths is not None:
+                speech_lengths = speech_lengths.squeeze(-1)
+            else:
+                speech_lengths = speech.shape[1]
+        else:
+            # extract fbank feats
+            time1 = time.perf_counter()
+            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,
+            )
+            time2 = time.perf_counter()
+            meta_data["load_data"] = f"{time2 - time1:0.3f}"
+            speech, speech_lengths = extract_fbank(
+                audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend
+            )
+            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
+            )
+
+        speech = speech.to(device=kwargs["device"])
+        speech_lengths = speech_lengths.to(device=kwargs["device"])
+        # Encoder
+        if kwargs.get("fp16", False):
+            speech = speech.half()
+        encoder_out, encoder_out_lens = self.encode(speech, speech_lengths)
+        if isinstance(encoder_out, tuple):
+            encoder_out = encoder_out[0]
+
+        ctc_probs = self.ctc.softmax(encoder_out)
+        ctc_greedy_paths = ctc_probs.argmax(dim=-1)
+
+        results = []
+        batch_size, n, d = encoder_out.size()
+        if isinstance(key[0], (list, tuple)):
+            key = key[0]
+        if len(key) < batch_size:
+            key = key * batch_size
+        for b in range(batch_size):
+
+            probs = ctc_probs[b, :encoder_out_lens[b]]
+            path = ctc_greedy_paths[b, :encoder_out_lens[b]]
+            
+            # Get compressed probabilities and timestamp indices
+            compressed_prob = self.average_repeats_inference(probs, path)
+            
+            # Handling Noise/Silence (Empty Output)
+            if compressed_prob.size(0) == 0:
+                token_int = []
+            else:
+                # 4. Decoder Forward
+                compressed_prob_in = compressed_prob.unsqueeze(0) # [1, U', V]
+                in_lens = torch.tensor([compressed_prob.size(0)], device=encoder_out.device)
+
+                decoder_out, _ = self.decoder(
+                        encoder_out[b:b+1],
+                        encoder_out_lens[b:b+1],
+                        compressed_prob_in,
+                        in_lens,)
+
+                
+                yseq = decoder_out.argmax(dim=-1)[0]
+                token_int = yseq.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
+                    )
+                )
+
+                result_i = {"key": key[b], "token_int": token_int}
+                results.append(result_i)
+
+        return results, meta_data
+
+    def export(self, **kwargs):
+        """Export.
+        
+            Args:
+                **kwargs: Additional keyword arguments.
+            """
+        from .export_meta import export_rebuild_model
+
+        if "max_seq_len" not in kwargs:
+            kwargs["max_seq_len"] = 512
+        models = export_rebuild_model(model=self, **kwargs)
+        return models