Initial commit: FunASR Speech Recognition Toolkit

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

funasr/__init__.py

@@ -0,0 +1,79 @@
+"""Initialize funasr package."""
+
+import importlib
+import os
+import pkgutil
+import traceback
+
+
+dirname = os.path.dirname(__file__)
+version_file = os.path.join(dirname, "version.txt")
+with open(version_file, "r") as f:
+    __version__ = f.read().strip()
+
+
+_IMPORT_ERRORS = {}
+_IMPORT_ERROR_TRACEBACKS = {}
+_IMPORT_DEBUG = os.environ.get("FUNASR_IMPORT_DEBUG") == "1"
+_STRICT_IMPORT = os.environ.get("FUNASR_STRICT_IMPORT") == "1"
+
+
+def _record_import_error(name, error):
+    """Internal: record import error.
+    
+        Args:
+            name: TODO.
+            error: TODO.
+        """
+    _IMPORT_ERRORS[name] = f"{error.__class__.__name__}: {error}"
+    _IMPORT_ERROR_TRACEBACKS[name] = traceback.format_exc()
+    if _IMPORT_DEBUG:
+        print(f"Failed to import {name}: {_IMPORT_ERRORS[name]}")
+
+
+def get_import_errors():
+    """Get import errors."""
+    return dict(_IMPORT_ERRORS)
+
+
+def get_import_error_tracebacks():
+    """Get import error tracebacks."""
+    return dict(_IMPORT_ERROR_TRACEBACKS)
+
+
+def import_submodules(package, recursive=True):
+    """Import submodules.
+    
+        Args:
+            package: TODO.
+            recursive: TODO.
+        """
+    if isinstance(package, str):
+        try:
+            package = importlib.import_module(package)
+        except Exception as e:
+            _record_import_error(package, e)
+            if _STRICT_IMPORT:
+                raise
+            return {}
+    results = {}
+    if not isinstance(package, str):
+        for loader, name, is_pkg in pkgutil.walk_packages(package.__path__, package.__name__ + "."):
+            try:
+                results[name] = importlib.import_module(name)
+            except Exception as e:
+                _record_import_error(name, e)
+                if _STRICT_IMPORT:
+                    raise
+                continue
+            if recursive and is_pkg:
+                results.update(import_submodules(name))
+    return results
+
+
+import_submodules(__name__)
+
+from funasr.auto.auto_model import AutoModel
+from funasr.auto.auto_frontend import AutoFrontend
+
+os.environ["HYDRA_FULL_ERROR"] = "1"

funasr/auto/auto_frontend.py

@@ -0,0 +1,122 @@
+#!/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 json
+import time
+import torch
+import hydra
+import random
+import string
+import logging
+import os.path
+from tqdm import tqdm
+from omegaconf import DictConfig, OmegaConf, ListConfig
+
+from funasr.register import tables
+from funasr.utils.load_utils import load_bytes
+from funasr.download.file import download_from_url
+from funasr.auto.auto_model import prepare_data_iterator
+from funasr.utils.timestamp_tools import timestamp_sentence
+from funasr.download.download_model_from_hub import download_model
+from funasr.utils.vad_utils import slice_padding_audio_samples
+from funasr.train_utils.set_all_random_seed import set_all_random_seed
+from funasr.train_utils.load_pretrained_model import load_pretrained_model
+from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
+from funasr.models.campplus.utils import sv_chunk, postprocess, distribute_spk
+
+
+class AutoFrontend:
+    def __init__(self, **kwargs):
+        """Initialize AutoFrontend.
+        
+            Args:
+                **kwargs: Additional keyword arguments.
+            """
+        assert "model" in kwargs
+        if "model_conf" not in kwargs:
+            logging.info("download models from model hub: {}".format(kwargs.get("hub", "ms")))
+            kwargs = download_model(**kwargs)
+
+        # build frontend
+        frontend = kwargs.get("frontend", None)
+        if frontend is not None:
+            frontend_class = tables.frontend_classes.get(frontend)
+            frontend = frontend_class(**kwargs["frontend_conf"])
+
+        self.frontend = frontend
+        if "frontend" in kwargs:
+            del kwargs["frontend"]
+        self.kwargs = kwargs
+
+    def __call__(self, input, input_len=None, kwargs=None, **cfg):
+
+        """Internal: call  .
+        
+            Args:
+                input: Input audio/text data.
+                input_len: TODO.
+                kwargs: Additional keyword arguments.
+                **cfg: Configuration overrides.
+            """
+        kwargs = self.kwargs if kwargs is None else kwargs
+        kwargs.update(cfg)
+
+        key_list, data_list = prepare_data_iterator(input, input_len=input_len)
+        batch_size = kwargs.get("batch_size", 1)
+        device = kwargs.get("device", "cuda")
+        if device == "cpu":
+            batch_size = 1
+
+        meta_data = {}
+
+        result_list = []
+        num_samples = len(data_list)
+        # pbar = tqdm(colour="blue", total=num_samples + 1, dynamic_ncols=True)
+
+        time0 = time.perf_counter()
+        for beg_idx in range(0, num_samples, batch_size):
+            end_idx = min(num_samples, beg_idx + batch_size)
+            data_batch = data_list[beg_idx:end_idx]
+            key_batch = key_list[beg_idx:end_idx]
+
+            # extract fbank feats
+            time1 = time.perf_counter()
+            audio_sample_list = load_audio_text_image_video(
+                data_batch, fs=self.frontend.fs, audio_fs=kwargs.get("fs", 16000)
+            )
+            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=self.frontend,
+                **kwargs,
+            )
+            time3 = time.perf_counter()
+            meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
+            meta_data["batch_data_time"] = (
+                speech_lengths.sum().item() * self.frontend.frame_shift * self.frontend.lfr_n / 1000
+            )
+
+            if kwargs.get("return_pt", True):
+                speech, speech_lengths = speech.to(device=device), speech_lengths.to(device=device)
+            else:
+                speech, speech_lengths = speech.numpy(), speech_lengths.numpy()
+            batch = {
+                "input": speech,
+                "input_len": speech_lengths,
+                "key": key_batch,
+                "data_type": "fbank",
+            }
+            result_list.append(batch)
+
+            # pbar.update(1)
+            # description = f"{meta_data}, "
+            # pbar.set_description(description)
+
+        time_end = time.perf_counter()
+        # pbar.set_description(f"time escaped total: {time_end - time0:0.3f}")
+
+        return result_list

funasr/auto/auto_model_vllm.py

@@ -0,0 +1,348 @@
+#!/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)
+
+"""
+Generic vLLM inference wrapper for ALL LLM-based ASR models in FunASR.
+
+Applicable models (any model with audio_encoder + adaptor + LLM architecture):
+    - FunASRNano (Fun-ASR-Nano-2512, Fun-ASR-MLT-Nano-2512)
+    - LLMASR (Whisper + Qwen/Vicuna/LLaMA)
+    - GLMASR (GLM-ASR-Nano)
+
+NOT applicable (these models don't use autoregressive LLM decoding):
+    - Paraformer (non-autoregressive CIF predictor + attention decoder)
+    - SenseVoice (Whisper-like encoder-decoder, not LLM-based)
+    - Conformer/Transformer ASR (CTC/attention, no LLM)
+    - CT-Transformer (punctuation model, small transformer)
+    - Qwen3-ASR (uses external qwen-asr package with its own optimized inference)
+
+Usage:
+    from funasr.auto.auto_model_vllm import AutoModelVLLM
+
+    # Works for any LLM-based ASR model
+    model = AutoModelVLLM(
+        model="FunAudioLLM/Fun-ASR-Nano-2512",
+        tensor_parallel_size=2,
+    )
+    results = model.generate(["audio.wav"])
+
+    # Also works for LLMASR models
+    model = AutoModelVLLM(
+        model="/path/to/llm_asr_model",
+        tensor_parallel_size=4,
+    )
+"""
+
+import glob
+import json
+import logging
+import os
+import re
+import shutil
+import time
+from typing import List, Optional, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+logger = logging.getLogger(__name__)
+
+dtype_map = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}
+
+# Models that use LLM and can benefit from vLLM
+_LLM_BASED_MODELS = {"FunASRNano", "LLMASR", "LLMASRNAR", "GLMASR", "QwenAudioWarp"}
+
+# Models that CANNOT use vLLM (no autoregressive LLM)
+_NON_LLM_MODELS = {
+    "Paraformer": "Non-autoregressive model using CIF predictor. No LLM decoding.",
+    "SenseVoice": "Whisper-like encoder-decoder. Not LLM-based.",
+    "CTTransformer": "Small punctuation model. No benefit from vLLM.",
+    "Conformer": "CTC/attention encoder-decoder. No LLM.",
+    "Qwen3ASR": "Uses external qwen-asr package with optimized inference.",
+}
+
+
+def check_vllm_applicable(model_name: str) -> bool:
+    """Check if a model can use vLLM inference.
+
+    Args:
+        model_name: The model class name from config.yaml.
+
+    Returns:
+        True if vLLM is applicable.
+
+    Raises:
+        ValueError: If model explicitly cannot use vLLM, with explanation.
+    """
+    if model_name in _LLM_BASED_MODELS:
+        return True
+    for non_llm, reason in _NON_LLM_MODELS.items():
+        if non_llm in model_name:
+            raise ValueError(
+                f"Model '{model_name}' cannot use vLLM: {reason}\n"
+                f"vLLM only accelerates autoregressive LLM decoding. "
+                f"Use the standard FunASR AutoModel for this model."
+            )
+    return False
+
+
+def prepare_vllm_weights(model_dir: str, output_dir: str = None) -> str:
+    """Extract LLM weights from model.pt into vLLM-compatible format.
+
+    Works for any model that stores LLM weights with 'llm.' prefix in model.pt
+    and has a config directory (e.g., Qwen3-0.6B/) with model config and tokenizer.
+
+    Args:
+        model_dir: Path to the FunASR model directory.
+        output_dir: Where to save extracted weights. Auto-detected if None.
+
+    Returns:
+        Path to vLLM-ready model directory.
+    """
+    if output_dir is None:
+        # Find the LLM config directory
+        from omegaconf import OmegaConf
+        config_path = os.path.join(model_dir, "config.yaml")
+        if os.path.exists(config_path):
+            config = OmegaConf.load(config_path)
+            llm_conf = OmegaConf.to_container(config.get("llm_conf", {}), resolve=True)
+            llm_path = llm_conf.get("init_param_path", "")
+            if llm_path and not os.path.isabs(llm_path):
+                llm_path = os.path.join(model_dir, llm_path)
+            if os.path.isdir(llm_path):
+                output_dir = llm_path + "-vllm"
+            else:
+                output_dir = os.path.join(model_dir, "llm-vllm")
+        else:
+            output_dir = os.path.join(model_dir, "llm-vllm")
+
+    # Check if already prepared
+    if glob.glob(os.path.join(output_dir, "*.safetensors")) or glob.glob(
+        os.path.join(output_dir, "model*.bin")
+    ):
+        logger.info(f"vLLM weights already at {output_dir}")
+        return output_dir
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Find and copy LLM config/tokenizer files
+    from omegaconf import OmegaConf
+    config = OmegaConf.load(os.path.join(model_dir, "config.yaml"))
+    llm_conf = OmegaConf.to_container(config.get("llm_conf", {}), resolve=True)
+    llm_config_dir = llm_conf.get("init_param_path", "")
+    if llm_config_dir and not os.path.isabs(llm_config_dir):
+        llm_config_dir = os.path.join(model_dir, llm_config_dir)
+
+    if os.path.isdir(llm_config_dir):
+        for fname in os.listdir(llm_config_dir):
+            src = os.path.join(llm_config_dir, fname)
+            dst = os.path.join(output_dir, fname)
+            if os.path.isfile(src) and not os.path.exists(dst):
+                shutil.copy2(src, dst)
+
+    # Extract LLM weights from model.pt
+    model_pt = os.path.join(model_dir, "model.pt")
+    if not os.path.exists(model_pt):
+        raise FileNotFoundError(f"model.pt not found at {model_pt}")
+
+    logger.info(f"Extracting LLM weights from {model_pt}...")
+    checkpoint = torch.load(model_pt, map_location="cpu")
+    state_dict = checkpoint.get("state_dict", checkpoint)
+
+    llm_state = {}
+    for key, value in state_dict.items():
+        if key.startswith("llm."):
+            llm_state[key[4:]] = value  # Remove 'llm.' prefix
+
+    if not llm_state:
+        raise RuntimeError("No LLM weights found (expected 'llm.*' prefix)")
+
+    logger.info(f"Extracted {len(llm_state)} LLM weight tensors")
+
+    try:
+        from safetensors.torch import save_file
+        save_path = os.path.join(output_dir, "model.safetensors")
+        save_file(llm_state, save_path)
+        index = {
+            "metadata": {"total_size": sum(v.numel() * v.element_size() for v in llm_state.values())},
+            "weight_map": {k: "model.safetensors" for k in llm_state.keys()},
+        }
+        with open(os.path.join(output_dir, "model.safetensors.index.json"), "w") as f:
+            json.dump(index, f, indent=2)
+    except ImportError:
+        torch.save(llm_state, os.path.join(output_dir, "model.bin"))
+
+    return output_dir
+
+
+class AutoModelVLLM:
+    """Generic vLLM wrapper for LLM-based ASR models.
+
+    Automatically detects model architecture, extracts LLM weights,
+    loads audio components in PyTorch, and uses vLLM for generation.
+
+    Works for: FunASRNano, LLMASR, GLMASR, and any model with
+    audio_encoder + audio_adaptor + LLM architecture.
+
+    Args:
+        model: Model name (hub) or local directory path.
+        hub: "ms" (ModelScope) or "hf" (HuggingFace).
+        device: Device for audio encoder/adaptor.
+        dtype: Compute dtype ("bf16", "fp16", "fp32").
+        tensor_parallel_size: GPUs for vLLM tensor parallelism.
+        gpu_memory_utilization: GPU memory fraction for vLLM.
+        max_model_len: Maximum sequence length.
+
+    Example:
+        >>> model = AutoModelVLLM(model="FunAudioLLM/Fun-ASR-Nano-2512")
+        >>> results = model.generate(["audio.wav"], language="中文")
+    """
+
+    def __init__(
+        self,
+        model: str,
+        hub: str = "ms",
+        device: str = "cuda:0",
+        dtype: str = "bf16",
+        tensor_parallel_size: int = 1,
+        gpu_memory_utilization: float = 0.8,
+        max_model_len: int = 4096,
+        enforce_eager: bool = False,
+        **kwargs,
+    ):
+        # Resolve model directory
+        if os.path.isdir(model):
+            self.model_dir = model
+        else:
+            if hub in ("ms", "modelscope"):
+                from modelscope.hub.snapshot_download import snapshot_download
+                self.model_dir = snapshot_download(model, revision=kwargs.get("revision", "master"))
+            elif hub in ("hf", "huggingface"):
+                from huggingface_hub import snapshot_download
+                self.model_dir = snapshot_download(model)
+            else:
+                raise ValueError(f"Unsupported hub: {hub}")
+
+        # Check model type
+        from omegaconf import OmegaConf
+        config = OmegaConf.load(os.path.join(self.model_dir, "config.yaml"))
+        self.model_type = config.get("model", "unknown")
+        check_vllm_applicable(self.model_type)
+
+        self.device = device
+        self.dtype = dtype
+        self.torch_dtype = dtype_map.get(dtype, torch.bfloat16)
+
+        # Use the specialized implementation if available
+        if self.model_type == "FunASRNano":
+            from funasr.models.fun_asr_nano.inference_vllm import FunASRNanoVLLM
+            self._engine = FunASRNanoVLLM(
+                model_dir=self.model_dir, device=device, dtype=dtype,
+                tensor_parallel_size=tensor_parallel_size,
+                gpu_memory_utilization=gpu_memory_utilization,
+                max_model_len=max_model_len, enforce_eager=enforce_eager,
+                **kwargs,
+            )
+        elif self.model_type in ("GLMASR", "glmasr"):
+            from funasr.models.glm_asr.inference_vllm import GLMASRVLLMEngine
+            self._engine = GLMASRVLLMEngine(
+                model_dir=self.model_dir, device=device, dtype=dtype,
+                tensor_parallel_size=tensor_parallel_size,
+                gpu_memory_utilization=gpu_memory_utilization,
+                max_model_len=max_model_len,
+                **kwargs,
+            )
+        elif self.model_type in ("LLMASR", "LLMASRNAR"):
+            self._engine = self._build_llmasr_engine(
+                config, tensor_parallel_size, gpu_memory_utilization,
+                max_model_len, enforce_eager, **kwargs,
+            )
+        else:
+            # Generic fallback using the FunASRNano approach
+            # (works for any model with audio_encoder + adaptor + LLM)
+            from funasr.models.fun_asr_nano.inference_vllm import FunASRNanoVLLM
+            self._engine = FunASRNanoVLLM(
+                model_dir=self.model_dir, device=device, dtype=dtype,
+                tensor_parallel_size=tensor_parallel_size,
+                gpu_memory_utilization=gpu_memory_utilization,
+                max_model_len=max_model_len, enforce_eager=enforce_eager,
+                **kwargs,
+            )
+
+    def _build_llmasr_engine(self, config, tensor_parallel_size, gpu_memory_utilization,
+                              max_model_len, enforce_eager, **kwargs):
+        """Build vLLM engine for LLMASR models (Whisper + Qwen/Vicuna)."""
+        from funasr.models.fun_asr_nano.inference_vllm import FunASRNanoVLLM
+
+        # LLMASR follows same pattern as FunASRNano
+        return FunASRNanoVLLM(
+            model_dir=self.model_dir, device=self.device, dtype=self.dtype,
+            tensor_parallel_size=tensor_parallel_size,
+            gpu_memory_utilization=gpu_memory_utilization,
+            max_model_len=max_model_len, enforce_eager=enforce_eager,
+            **kwargs,
+        )
+
+    def generate(self, inputs, **kwargs):
+        """Run ASR inference.
+
+        Args:
+            inputs: Audio file path(s), numpy arrays, or tensors.
+            **kwargs: Model-specific parameters (language, hotwords, etc.)
+
+        Returns:
+            List of result dicts with "key" and "text" fields.
+        """
+        self._warn_if_audio_too_long(inputs)
+        return self._engine.generate(inputs, **kwargs)
+
+    def _warn_if_audio_too_long(self, inputs, max_safe_sec=40.0):
+        """Warn (once) if a single audio input is long enough to be truncated.
+
+        Fun-ASR-Nano is a segment-level (LLM-)ASR model. Decoding very long
+        audio in a single pass can silently truncate or degrade the output -- the
+        decode hits ``max_new_tokens`` long before the audio ends, so the user
+        gets a partial transcript with no error. The right usage is to
+        pre-segment with VAD; this warning points users there instead of letting
+        them get a silently truncated result. It does not change the output.
+        """
+        if getattr(self, "_warned_audio_too_long", False):
+            return
+        items = inputs if isinstance(inputs, (list, tuple)) else [inputs]
+        for item in items:
+            duration = None
+            try:
+                if isinstance(item, str):
+                    import soundfile as sf
+
+                    duration = sf.info(item).duration
+                elif isinstance(item, np.ndarray):
+                    duration = item.shape[-1] / 16000.0
+                elif isinstance(item, torch.Tensor):
+                    duration = item.shape[-1] / 16000.0
+            except Exception:
+                continue
+            if duration is not None and duration > max_safe_sec:
+                logger.warning(
+                    "AutoModelVLLM received a %.0fs audio input. Fun-ASR-Nano is a "
+                    "segment-level model; decoding very long audio in a single pass can "
+                    "truncate or degrade the result. Pre-segment with VAD and pass the "
+                    "segments, or use the high-level `funasr.AutoModel(model=..., "
+                    'vad_model="fsmn-vad")`, which segments long audio automatically.',
+                    duration,
+                )
+                self._warned_audio_too_long = True
+                break
+
+    @classmethod
+    def supported_models(cls):
+        """Return dict of model types and their vLLM support status."""
+        info = {}
+        for m in _LLM_BASED_MODELS:
+            info[m] = {"supported": True, "reason": "LLM-based, autoregressive generation"}
+        for m, reason in _NON_LLM_MODELS.items():
+            info[m] = {"supported": False, "reason": reason}
+        return info

funasr/auto/auto_tokenizer.py

@@ -0,0 +1,8 @@
+class AutoTokenizer:
+    """
+    Undo
+    """
+
+    def __init__(self):
+        """Initialize AutoTokenizer."""
+        pass

funasr/bin/_server_app.py

@@ -0,0 +1,298 @@
+"""FunASR Server — unified vLLM-based inference service.
+
+Provides OpenAI-compatible API (/v1/audio/transcriptions) and REST API (/asr).
+Uses vLLM for Fun-ASR-Nano (GPU) or falls back to AutoModel for non-LLM models (SenseVoice/Paraformer).
+"""
+
+import io
+import os
+import re
+import time
+import logging
+import tempfile
+from typing import Optional
+
+import numpy as np
+import soundfile as sf
+
+try:
+    from fastapi import FastAPI, UploadFile, File, Form, HTTPException
+    from fastapi.responses import JSONResponse
+except ImportError:
+    raise ImportError(
+        "funasr-server requires additional packages. Install with: pip install vllm fastapi uvicorn python-multipart"
+    )
+
+logger = logging.getLogger("funasr.server")
+
+
+def prepare_audio_for_inference(audio_data, sr, target_sr=16000):
+    """Return mono float32 audio at target_sr for ASR inference."""
+    audio_data = np.asarray(audio_data)
+    if audio_data.ndim > 1:
+        channel_axis = -1 if audio_data.shape[-1] <= audio_data.shape[0] else 0
+        audio_data = audio_data.mean(axis=channel_axis)
+
+    if sr != target_sr:
+        import librosa
+        audio_data = librosa.resample(audio_data, orig_sr=sr, target_sr=target_sr)
+        sr = target_sr
+
+    return audio_data.astype(np.float32), sr
+
+def create_app(device: str = "cuda", preload_model: str = "auto") -> FastAPI:
+    if preload_model == "auto":
+        preload_model = "fun-asr-nano" if device.startswith("cuda") else "sensevoice"
+
+    app = FastAPI(title="FunASR Server", version="1.3.6")
+    app.state.device = device
+    app.state.engine = None
+    app.state.vad_model = None
+    app.state.fallback_models = {}
+
+    # Non-LLM model configs (use AutoModel, no vLLM)
+    FALLBACK_CONFIGS = {
+        "sensevoice": {
+            "model": "iic/SenseVoiceSmall",
+            "vad_model": "fsmn-vad",
+            "vad_kwargs": {"max_single_segment_time": 30000},
+        },
+        "paraformer": {
+            "model": "paraformer-zh",
+            "vad_model": "fsmn-vad",
+            "punc_model": "ct-punc",
+        },
+    }
+
+    def _load_vllm_engine():
+        """Load Fun-ASR-Nano vLLM engine. Falls back to AutoModel if vLLM unavailable."""
+        if app.state.engine is not None:
+            return
+        try:
+            from funasr.models.fun_asr_nano.inference_vllm import FunASRNanoVLLM
+            from funasr import AutoModel as _AutoModel
+
+            logger.info("Loading Fun-ASR-Nano vLLM engine...")
+            t0 = time.time()
+            app.state.engine = FunASRNanoVLLM.from_pretrained(
+                model="FunAudioLLM/Fun-ASR-Nano-2512",
+                hub="hf",
+                device=device,
+                dtype="bf16",
+                max_model_len=4096,
+                gpu_memory_utilization=0.5,
+            )
+            logger.info(f"vLLM engine ready in {time.time()-t0:.1f}s")
+            app.state.use_vllm = True
+
+            logger.info("Loading VAD model...")
+            app.state.vad_model = _AutoModel(model="fsmn-vad", device=device, disable_update=True)
+            logger.info("VAD ready.")
+        except Exception as e:
+            logger.warning(f"vLLM failed ({e}), falling back to AutoModel for fun-asr-nano")
+            app.state.use_vllm = False
+            from funasr import AutoModel
+            cfg = {
+                "model": "FunAudioLLM/Fun-ASR-Nano-2512",
+                "hub": "hf",
+                "trust_remote_code": True,
+                "vad_model": "fsmn-vad",
+                "vad_kwargs": {"max_single_segment_time": 30000},
+                "device": device,
+                "disable_update": True,
+            }
+            app.state.fallback_models["fun-asr-nano"] = AutoModel(**cfg)
+            logger.info("Fallback AutoModel loaded for fun-asr-nano.")
+
+    def _load_fallback(name: str):
+        """Load non-LLM model via AutoModel."""
+        if name in app.state.fallback_models:
+            return app.state.fallback_models[name]
+        if name not in FALLBACK_CONFIGS:
+            return None
+        from funasr import AutoModel
+        cfg = FALLBACK_CONFIGS[name].copy()
+        cfg["device"] = device
+        cfg["disable_update"] = True
+        logger.info(f"Loading fallback model '{name}'...")
+        model = AutoModel(**cfg)
+        app.state.fallback_models[name] = model
+        return model
+
+    def _process_vllm(audio_data, sr, language=None, hotwords=None, use_spk=False):
+        """Process audio with vLLM engine (Fun-ASR-Nano)."""
+        audio_data, sr = prepare_audio_for_inference(audio_data, sr)
+
+        # VAD
+        vad_res = app.state.vad_model.generate(input=audio_data, fs=sr)
+        segments = vad_res[0]["value"] if vad_res and vad_res[0].get("value") else [[0, int(len(audio_data)*1000/sr)]]
+
+        seg_audios = []
+        seg_times = []
+        for seg in segments:
+            s0 = int(seg[0] * sr / 1000)
+            s1 = int(seg[1] * sr / 1000)
+            seg_audio = audio_data[s0:s1]
+            if len(seg_audio) > sr * 0.3:
+                seg_audios.append(seg_audio)
+                seg_times.append((seg[0], seg[1]))
+
+        if not seg_audios:
+            return {"text": "", "segments": [], "duration": len(audio_data)/sr}
+
+        # repetition_penalty is left at the neutral 1.0: the Fun-ASR-Nano vLLM
+        # engine runs in prompt-embeds mode, where any other value crashes the
+        # CUDA kernel (see issue #2948 and fun_asr_nano.vllm_utils).
+        gen_kwargs = {"max_new_tokens": 500, "repetition_penalty": 1.0}
+        if language:
+            gen_kwargs["language"] = language
+        if hotwords:
+            gen_kwargs["hotwords"] = hotwords
+
+        results = app.state.engine.generate(inputs=seg_audios, **gen_kwargs)
+
+        output_segments = []
+        full_text_parts = []
+        for r, (start_ms, end_ms) in zip(results, seg_times):
+            text = r["text"]
+            seg_info = {"text": text, "start": start_ms/1000, "end": end_ms/1000}
+            if "timestamps" in r:
+                offset = start_ms / 1000
+                seg_info["words"] = [
+                    {"word": ts["token"], "start": ts["start_time"]+offset, "end": ts["end_time"]+offset}
+                    for ts in r["timestamps"]
+                ]
+            output_segments.append(seg_info)
+            full_text_parts.append(text)
+
+        return {
+            "text": "".join(full_text_parts),
+            "segments": output_segments,
+            "duration": len(audio_data) / sr,
+        }
+
+    def _process_fallback(model_name, audio_path, language=None):
+        """Process with non-LLM model (SenseVoice/Paraformer)."""
+        model = _load_fallback(model_name)
+        kwargs = {"input": audio_path, "batch_size": 1}
+        if language:
+            kwargs["language"] = language
+        result = model.generate(**kwargs)
+        text = re.sub(r'<\|[^|]*\|>', '', result[0]["text"]).strip()
+        segments = []
+        if "sentence_info" in result[0]:
+            for s in result[0]["sentence_info"]:
+                segments.append({
+                    "start": s.get("start", 0)/1000,
+                    "end": s.get("end", 0)/1000,
+                    "text": re.sub(r'<\|[^|]*\|>', '', s.get("text", "")).strip(),
+                    "speaker": s.get("spk"),
+                })
+        return {"text": text, "segments": segments}
+
+    # Pre-load
+    if preload_model == "fun-asr-nano":
+        _load_vllm_engine()
+    else:
+        _load_fallback(preload_model)
+
+    @app.post("/v1/audio/transcriptions")
+    async def transcribe(
+        file: UploadFile = File(...),
+        model: str = Form(default="fun-asr-nano"),
+        language: Optional[str] = Form(default=None),
+        response_format: Optional[str] = Form(default="json"),
+        spk: bool = Form(default=False),
+    ):
+        content = await file.read()
+        t0 = time.perf_counter()
+
+        if model == "fun-asr-nano":
+            _load_vllm_engine()
+            if app.state.use_vllm:
+                audio_data, sr = sf.read(io.BytesIO(content))
+                result = _process_vllm(audio_data, sr, language=language, use_spk=spk)
+            else:
+                suffix = os.path.splitext(file.filename)[1] if file.filename else ".wav"
+                with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
+                    tmp.write(content)
+                    tmp_path = tmp.name
+                try:
+                    result = _process_fallback("fun-asr-nano", tmp_path, language=language)
+                finally:
+                    os.unlink(tmp_path)
+        elif model in FALLBACK_CONFIGS:
+            suffix = os.path.splitext(file.filename)[1] if file.filename else ".wav"
+            with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
+                tmp.write(content)
+                tmp_path = tmp.name
+            try:
+                result = _process_fallback(model, tmp_path, language=language)
+            finally:
+                os.unlink(tmp_path)
+        else:
+            raise HTTPException(400, f"Unknown model '{model}'. Available: fun-asr-nano, {', '.join(FALLBACK_CONFIGS.keys())}")
+
+        t1 = time.perf_counter()
+
+        if response_format == "verbose_json":
+            return JSONResponse({
+                "task": "transcribe",
+                "language": language or "zh",
+                "duration": result.get("duration", 0),
+                "text": result["text"],
+                "segments": [
+                    {"id": i, "start": s["start"], "end": s["end"], "text": s["text"], "words": s.get("words", [])}
+                    for i, s in enumerate(result["segments"])
+                ],
+            })
+        elif response_format == "text":
+            return JSONResponse(result["text"])
+        else:
+            return JSONResponse({"text": result["text"]})
+
+    @app.post("/asr")
+    async def asr_endpoint(
+        file: UploadFile = File(...),
+        language: Optional[str] = Form(default=None),
+        hotwords: str = Form(default=""),
+        spk: bool = Form(default=False),
+    ):
+        """Full-featured ASR endpoint with timestamps and speaker diarization."""
+        content = await file.read()
+        _load_vllm_engine()
+        hw_list = [w.strip() for w in hotwords.split(",") if w.strip()] if hotwords else None
+
+        t0 = time.perf_counter()
+        if app.state.use_vllm:
+            audio_data, sr = sf.read(io.BytesIO(content))
+            result = _process_vllm(audio_data, sr, language=language, hotwords=hw_list, use_spk=spk)
+        else:
+            suffix = ".wav"
+            with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
+                tmp.write(content)
+                tmp_path = tmp.name
+            try:
+                result = _process_fallback("fun-asr-nano", tmp_path, language=language)
+            finally:
+                os.unlink(tmp_path)
+        t1 = time.perf_counter()
+
+        result["processing_time"] = round(t1 - t0, 3)
+        result["rtf"] = round((t1 - t0) / result["duration"], 4) if result.get("duration", 0) > 0 else 0
+        return JSONResponse(result)
+
+    @app.get("/v1/models")
+    async def list_models():
+        all_models = ["fun-asr-nano"] + list(FALLBACK_CONFIGS.keys())
+        return JSONResponse({"object": "list", "data": [{"id": n, "object": "model"} for n in all_models]})
+
+    @app.get("/health")
+    async def health():
+        loaded = []
+        if app.state.engine is not None:
+            loaded.append("fun-asr-nano (vLLM)")
+        loaded.extend(app.state.fallback_models.keys())
+        return {"status": "ok", "device": device, "models_loaded": loaded}
+
+    return app

funasr/bin/compute_audio_cmvn.py

@@ -0,0 +1,146 @@
+import os
+import json
+import numpy as np
+import torch
+import hydra
+import logging
+from omegaconf import DictConfig, OmegaConf
+
+from funasr.register import tables
+from funasr.download.download_model_from_hub import download_model
+from funasr.train_utils.set_all_random_seed import set_all_random_seed
+
+
+@hydra.main(config_name=None, version_base=None)
+def main_hydra(kwargs: DictConfig):
+    """Main hydra.
+    
+        Args:
+            kwargs: Additional keyword arguments.
+        """
+    if kwargs.get("debug", False):
+        import pdb
+
+        pdb.set_trace()
+
+    assert "model" in kwargs
+    if "model_conf" not in kwargs:
+        logging.info("download models from model hub: {}".format(kwargs.get("hub", "ms")))
+        kwargs = download_model(is_training=kwargs.get("is_training", True), **kwargs)
+
+    main(**kwargs)
+
+
+def main(**kwargs):
+    """Main.
+    
+        Args:
+            **kwargs: Additional keyword arguments.
+        """
+    print(kwargs)
+    # set random seed
+    # tables.print()
+    set_all_random_seed(kwargs.get("seed", 0))
+    torch.backends.cudnn.enabled = kwargs.get("cudnn_enabled", torch.backends.cudnn.enabled)
+    torch.backends.cudnn.benchmark = kwargs.get("cudnn_benchmark", torch.backends.cudnn.benchmark)
+    torch.backends.cudnn.deterministic = kwargs.get("cudnn_deterministic", True)
+
+    tokenizer = kwargs.get("tokenizer", None)
+
+    # build frontend if frontend is none None
+    frontend = kwargs.get("frontend", None)
+    if frontend is not None:
+        frontend_class = tables.frontend_classes.get(frontend)
+        frontend = frontend_class(**kwargs["frontend_conf"])
+        kwargs["frontend"] = frontend
+        kwargs["input_size"] = frontend.output_size()
+
+    # dataset
+    dataset_class = tables.dataset_classes.get(kwargs.get("dataset", "AudioDataset"))
+    dataset_train = dataset_class(
+        kwargs.get("train_data_set_list"),
+        frontend=frontend,
+        tokenizer=None,
+        is_training=False,
+        **kwargs.get("dataset_conf"),
+    )
+
+    # dataloader
+    batch_sampler = kwargs["dataset_conf"].get("batch_sampler", "BatchSampler")
+    batch_sampler_class = tables.batch_sampler_classes.get(batch_sampler)
+    dataset_conf = kwargs.get("dataset_conf")
+    dataset_conf["batch_type"] = "example"
+    dataset_conf["batch_size"] = 1
+    dataset_conf["num_workers"] = os.cpu_count() or 32
+    batch_sampler_train = batch_sampler_class(dataset_train, is_training=False, **dataset_conf)
+
+    dataloader_train = torch.utils.data.DataLoader(
+        dataset_train, collate_fn=dataset_train.collator, **batch_sampler_train
+    )
+
+    total_frames = 0
+    for batch_idx, batch in enumerate(dataloader_train):
+        iter_stop = int(kwargs.get("scale", -1.0) * len(dataloader_train))
+        log_step = iter_stop // 100
+        if batch_idx % log_step == 0:
+            logging.info(f"prcessed: {batch_idx}/{iter_stop}")
+        if batch_idx >= iter_stop and iter_stop > 0.0:
+            logging.info(f"prcessed: {iter_stop}/{iter_stop}")
+            break
+
+        fbank = batch["speech"].numpy()[0, :, :]
+        if total_frames == 0:
+            mean_stats = np.sum(fbank, axis=0)
+            var_stats = np.sum(np.square(fbank), axis=0)
+        else:
+            mean_stats += np.sum(fbank, axis=0)
+            var_stats += np.sum(np.square(fbank), axis=0)
+        total_frames += fbank.shape[0]
+
+    cmvn_info = {
+        "mean_stats": mean_stats.tolist(),
+        "var_stats": var_stats.tolist(),
+        "total_frames": total_frames,
+    }
+    cmvn_file = kwargs.get("cmvn_file", "cmvn.json")
+    # import pdb;pdb.set_trace()
+    with open(cmvn_file, "w") as fout:
+        fout.write(json.dumps(cmvn_info))
+
+    mean = -1.0 * mean_stats / total_frames
+    var = 1.0 / np.sqrt(var_stats / total_frames - mean * mean)
+    dims = mean.shape[0]
+    am_mvn = os.path.dirname(cmvn_file) + "/am.mvn"
+    with open(am_mvn, "w") as fout:
+        fout.write(
+            "<Nnet>"
+            + "\n"
+            + "<Splice> "
+            + str(dims)
+            + " "
+            + str(dims)
+            + "\n"
+            + "[ 0 ]"
+            + "\n"
+            + "<AddShift> "
+            + str(dims)
+            + " "
+            + str(dims)
+            + "\n"
+        )
+        fout.write("<LearnRateCoef> 0 [ " + " ".join([str(item) for item in mean]) + " ]\n")
+        fout.write("<Rescale> " + str(dims) + " " + str(dims) + "\n")
+        fout.write("<LearnRateCoef> 0 [ " + " ".join([str(item) for item in var]) + " ]\n")
+        fout.write("</Nnet>" + "\n")
+
+
+"""
+python funasr/bin/compute_audio_cmvn.py \
+--config-path "/Users/zhifu/funasr1.0/examples/aishell/paraformer/conf" \
+--config-name "train_asr_paraformer_conformer_12e_6d_2048_256.yaml" \
+++train_data_set_list="/Users/zhifu/funasr1.0/data/list/audio_datasets.jsonl" \
+++cmvn_file="/Users/zhifu/funasr1.0/data/list/cmvn.json" \
+++dataset_conf.num_workers=0
+"""
+if __name__ == "__main__":
+    main_hydra()

funasr/bin/export.py

@@ -0,0 +1,52 @@
+import os
+import hydra
+import logging
+from omegaconf import DictConfig, OmegaConf, ListConfig
+
+from funasr.auto.auto_model import AutoModel
+
+
+@hydra.main(config_name=None, version_base=None)
+def main_hydra(cfg: DictConfig):
+    """Main hydra.
+    
+        Args:
+            cfg: Configuration overrides.
+        """
+    def to_plain_list(cfg_item):
+        """To plain list.
+        
+            Args:
+                cfg_item: TODO.
+            """
+        if isinstance(cfg_item, ListConfig):
+            return OmegaConf.to_container(cfg_item, resolve=True)
+        elif isinstance(cfg_item, DictConfig):
+            return {k: to_plain_list(v) for k, v in cfg_item.items()}
+        else:
+            return cfg_item
+
+    kwargs = to_plain_list(cfg)
+
+    if kwargs.get("debug", False):
+        import pdb
+
+        pdb.set_trace()
+
+    if "device" not in kwargs:
+        kwargs["device"] = "cpu"
+    model = AutoModel(**kwargs)
+
+    res = model.export(
+        input=kwargs.get("input", None),
+        type=kwargs.get("type", "onnx"),
+        quantize=kwargs.get("quantize", False),
+        fallback_num=kwargs.get("fallback-num", 5),
+        calib_num=kwargs.get("calib_num", 100),
+        opset_version=kwargs.get("opset_version", 14),
+    )
+    print(res)
+
+
+if __name__ == "__main__":
+    main_hydra()

funasr/bin/inference.py

@@ -0,0 +1,40 @@
+import hydra
+import logging
+from omegaconf import DictConfig, OmegaConf, ListConfig
+
+from funasr.auto.auto_model import AutoModel
+
+
+@hydra.main(config_name=None, version_base=None)
+def main_hydra(cfg: DictConfig):
+    """Main hydra.
+    
+        Args:
+            cfg: Configuration overrides.
+        """
+    def to_plain_list(cfg_item):
+        """To plain list.
+        
+            Args:
+                cfg_item: TODO.
+            """
+        if isinstance(cfg_item, ListConfig):
+            return OmegaConf.to_container(cfg_item, resolve=True)
+        elif isinstance(cfg_item, DictConfig):
+            return {k: to_plain_list(v) for k, v in cfg_item.items()}
+        else:
+            return cfg_item
+
+    kwargs = to_plain_list(cfg)
+
+    if kwargs.get("debug", False):
+        import pdb
+
+        pdb.set_trace()
+    model = AutoModel(**kwargs)
+    res = model.generate(input=kwargs["input"])
+    print(res)
+
+
+if __name__ == "__main__":
+    main_hydra()

funasr/bin/server.py

@@ -0,0 +1,66 @@
+"""
+FunASR Server — OpenAI-compatible speech recognition API.
+
+Usage:
+    funasr-server                          # default: sensevoice on cuda:0, port 8000
+    funasr-server --device cpu --port 9000
+    funasr-server --model paraformer
+"""
+
+import argparse
+import sys
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="FunASR OpenAI-Compatible API Server",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  funasr-server                          # Start with SenseVoice on GPU
+  funasr-server --device cpu             # Start on CPU
+  funasr-server --model paraformer       # Use Paraformer model
+  funasr-server --port 9000             # Custom port
+
+Then use with OpenAI SDK:
+  from openai import OpenAI
+  client = OpenAI(base_url="http://localhost:8000/v1", api_key="x")
+  result = client.audio.transcriptions.create(model="sensevoice", file=open("a.wav","rb"))
+""",
+    )
+    parser.add_argument("--host", default="0.0.0.0", help="Bind address (default: 0.0.0.0)")
+    parser.add_argument("--port", type=int, default=8000, help="Port (default: 8000)")
+    parser.add_argument("--device", default="cuda", help="Device: cuda, cpu, mps (default: cuda)")
+    parser.add_argument("--model", default="auto", help="Pre-load model: auto (GPU=fun-asr-nano, CPU=sensevoice), sensevoice, paraformer, fun-asr-nano")
+    args = parser.parse_args()
+
+    try:
+        import uvicorn
+        import fastapi
+    except ImportError:
+        print("Error: funasr-server requires additional packages.")
+        print("Install with: pip install vllm fastapi uvicorn python-multipart")
+        sys.exit(1)
+
+    # Import and configure the app
+    import os
+    sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'examples', 'openai_api'))
+
+    # Use inline app to avoid path issues
+    from funasr.bin._server_app import create_app
+
+    app = create_app(device=args.device, preload_model=args.model)
+
+    print(f"╔══════════════════════════════════════════════╗")
+    print(f"║  FunASR Server v1.3.6                        ║")
+    print(f"║  Device: {args.device:<8}                          ║")
+    print(f"║  Model:  {args.model:<12}                      ║")
+    print(f"║  URL:    http://{args.host}:{args.port}/v1          ║")
+    print(f"║  Docs:   http://{args.host}:{args.port}/docs        ║")
+    print(f"╚══════════════════════════════════════════════╝")
+
+    uvicorn.run(app, host=args.host, port=args.port)
+
+
+if __name__ == "__main__":
+    main()

funasr/bin/tokenize_text.py

@@ -0,0 +1,307 @@
+#!/usr/bin/env python3
+import argparse
+from collections import Counter
+import logging
+from pathlib import Path
+import sys
+from typing import List
+from typing import Optional
+
+
+try:
+    from funasr.utils.cli_utils import get_commandline_args
+except ImportError:
+    def get_commandline_args():
+        return {}
+from funasr.tokenizer.build_tokenizer import build_tokenizer
+from funasr.tokenizer.cleaner import TextCleaner
+from funasr.tokenizer.phoneme_tokenizer import g2p_classes
+from funasr.utils.types import str2bool
+from funasr.utils.types import str_or_none
+
+
+def field2slice(field: Optional[str]) -> slice:
+    """Convert field string to slice
+
+    Note that field string accepts 1-based integer.
+
+    Examples:
+        >>> field2slice("1-")
+        slice(0, None, None)
+        >>> field2slice("1-3")
+        slice(0, 3, None)
+        >>> field2slice("-3")
+        slice(None, 3, None)
+    """
+    field = field.strip()
+    try:
+        if "-" in field:
+            # e.g. "2-" or "2-5" or "-7"
+            s1, s2 = field.split("-", maxsplit=1)
+            if s1.strip() == "":
+                s1 = None
+            else:
+                s1 = int(s1)
+                if s1 == 0:
+                    raise ValueError("1-based string")
+            if s2.strip() == "":
+                s2 = None
+            else:
+                s2 = int(s2)
+        else:
+            # e.g. "2"
+            s1 = int(field)
+            s2 = s1 + 1
+            if s1 == 0:
+                raise ValueError("must be 1 or more value")
+    except ValueError:
+        raise RuntimeError(f"Format error: e.g. '2-', '2-5', or '-5': {field}")
+
+    if s1 is None:
+        slic = slice(None, s2)
+    else:
+        # -1 because of 1-based integer following "cut" command
+        # e.g "1-3" -> slice(0, 3)
+        slic = slice(s1 - 1, s2)
+    return slic
+
+
+def tokenize(
+    input: str,
+    output: str,
+    field: Optional[str],
+    delimiter: Optional[str],
+    token_type: str,
+    space_symbol: str,
+    non_linguistic_symbols: Optional[str],
+    bpemodel: Optional[str],
+    log_level: str,
+    write_vocabulary: bool,
+    vocabulary_size: int,
+    remove_non_linguistic_symbols: bool,
+    cutoff: int,
+    add_symbol: List[str],
+    cleaner: Optional[str],
+    g2p: Optional[str],
+):
+
+    """Tokenize.
+    
+        Args:
+            input: Input audio/text data.
+            output: TODO.
+            field: TODO.
+            delimiter: TODO.
+            token_type: TODO.
+            space_symbol: TODO.
+            non_linguistic_symbols: TODO.
+            bpemodel: TODO.
+            log_level: TODO.
+            write_vocabulary: TODO.
+            vocabulary_size: Size/dimension parameter.
+            remove_non_linguistic_symbols: TODO.
+            cutoff: TODO.
+            add_symbol: TODO.
+            cleaner: TODO.
+            g2p: TODO.
+        """
+    logging.basicConfig(
+        level=log_level,
+        format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
+    )
+    if input == "-":
+        fin = sys.stdin
+    else:
+        fin = Path(input).open("r", encoding="utf-8")
+    if output == "-":
+        fout = sys.stdout
+    else:
+        p = Path(output)
+        p.parent.mkdir(parents=True, exist_ok=True)
+        fout = p.open("w", encoding="utf-8")
+
+    cleaner = TextCleaner(cleaner)
+    tokenizer = build_tokenizer(
+        token_type=token_type,
+        bpemodel=bpemodel,
+        delimiter=delimiter,
+        space_symbol=space_symbol,
+        non_linguistic_symbols=non_linguistic_symbols,
+        remove_non_linguistic_symbols=remove_non_linguistic_symbols,
+        g2p_type=g2p,
+    )
+
+    counter = Counter()
+    if field is not None:
+        field = field2slice(field)
+
+    for line in fin:
+        line = line.rstrip()
+        if field is not None:
+            # e.g. field="2-"
+            # uttidA hello world!! -> hello world!!
+            tokens = line.split(delimiter)
+            tokens = tokens[field]
+            if delimiter is None:
+                line = " ".join(tokens)
+            else:
+                line = delimiter.join(tokens)
+
+        line = cleaner(line)
+        tokens = tokenizer.text2tokens(line)
+        if not write_vocabulary:
+            fout.write(" ".join(tokens) + "\n")
+        else:
+            for t in tokens:
+                counter[t] += 1
+
+    if not write_vocabulary:
+        return
+
+    ## FIXME
+    ## del duplicate add_symbols in counter
+    for symbol_and_id in add_symbol:
+        # e.g symbol="<blank>:0"
+        try:
+            symbol, idx = symbol_and_id.split(":")
+        except ValueError:
+            raise RuntimeError(f"Format error: e.g. '<blank>:0': {symbol_and_id}")
+        symbol = symbol.strip()
+        if symbol in counter:
+            del counter[symbol]
+
+    # ======= write_vocabulary mode from here =======
+    # Sort by the number of occurrences in descending order
+    # and filter lower frequency words than cutoff value
+    words_and_counts = list(
+        filter(lambda x: x[1] > cutoff, sorted(counter.items(), key=lambda x: -x[1]))
+    )
+    # Restrict the vocabulary size
+    if vocabulary_size > 0:
+        if vocabulary_size < len(add_symbol):
+            raise RuntimeError(f"vocabulary_size is too small: {vocabulary_size}")
+        words_and_counts = words_and_counts[: vocabulary_size - len(add_symbol)]
+
+    # Parse the values of --add_symbol
+    for symbol_and_id in add_symbol:
+        # e.g symbol="<blank>:0"
+        try:
+            symbol, idx = symbol_and_id.split(":")
+            idx = int(idx)
+        except ValueError:
+            raise RuntimeError(f"Format error: e.g. '<blank>:0': {symbol_and_id}")
+        symbol = symbol.strip()
+
+        # e.g. idx=0  -> append as the first symbol
+        # e.g. idx=-1 -> append as the last symbol
+        if idx < 0:
+            idx = len(words_and_counts) + 1 + idx
+        words_and_counts.insert(idx, (symbol, None))
+
+    # Write words
+    for w, c in words_and_counts:
+        fout.write(w + "\n")
+
+    # Logging
+    total_count = sum(counter.values())
+    invocab_count = sum(c for w, c in words_and_counts if c is not None)
+    logging.info(f"OOV rate = {(total_count - invocab_count) / total_count * 100} %")
+
+
+def get_parser() -> argparse.ArgumentParser:
+    """Get parser."""
+    parser = argparse.ArgumentParser(
+        description="Tokenize texts",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "--log_level",
+        type=lambda x: x.upper(),
+        default="INFO",
+        choices=("CRITICAL", "ERROR", "WARNING", "INFO", "DEBUG", "NOTSET"),
+        help="The verbose level of logging",
+    )
+
+    parser.add_argument("--input", "-i", required=True, help="Input text. - indicates sys.stdin")
+    parser.add_argument("--output", "-o", required=True, help="Output text. - indicates sys.stdout")
+    parser.add_argument(
+        "--field",
+        "-f",
+        help="The target columns of the input text as 1-based integer. e.g 2-",
+    )
+    parser.add_argument(
+        "--token_type",
+        "-t",
+        default="char",
+        choices=["char", "bpe", "word", "phn"],
+        help="Token type",
+    )
+    parser.add_argument("--delimiter", "-d", default=None, help="The delimiter")
+    parser.add_argument("--space_symbol", default="<space>", help="The space symbol")
+    parser.add_argument("--bpemodel", default=None, help="The bpemodel file path")
+    parser.add_argument(
+        "--non_linguistic_symbols",
+        type=str_or_none,
+        help="non_linguistic_symbols file path",
+    )
+    parser.add_argument(
+        "--remove_non_linguistic_symbols",
+        type=str2bool,
+        default=False,
+        help="Remove non-language-symbols from tokens",
+    )
+    parser.add_argument(
+        "--cleaner",
+        type=str_or_none,
+        choices=[None, "tacotron", "jaconv", "vietnamese", "korean_cleaner"],
+        default=None,
+        help="Apply text cleaning",
+    )
+    parser.add_argument(
+        "--g2p",
+        type=str_or_none,
+        choices=g2p_classes,
+        default=None,
+        help="Specify g2p method if --token_type=phn",
+    )
+
+    group = parser.add_argument_group("write_vocabulary mode related")
+    group.add_argument(
+        "--write_vocabulary",
+        type=str2bool,
+        default=False,
+        help="Write tokens list instead of tokenized text per line",
+    )
+    group.add_argument("--vocabulary_size", type=int, default=0, help="Vocabulary size")
+    group.add_argument(
+        "--cutoff",
+        default=0,
+        type=int,
+        help="cut-off frequency used for write-vocabulary mode",
+    )
+    group.add_argument(
+        "--add_symbol",
+        type=str,
+        default=[],
+        action="append",
+        help="Append symbol e.g. --add_symbol '<blank>:0' --add_symbol '<unk>:1'",
+    )
+
+    return parser
+
+
+def main(cmd=None):
+    """Main.
+    
+        Args:
+            cmd: TODO.
+        """
+    print(get_commandline_args(), file=sys.stderr)
+    parser = get_parser()
+    args = parser.parse_args(cmd)
+    kwargs = vars(args)
+    tokenize(**kwargs)
+
+
+if __name__ == "__main__":
+    main()

funasr/bin/train.py

@@ -0,0 +1,289 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+
+import os
+import functools
+import sys
+import torch
+import torch.nn as nn
+import hydra
+import logging
+import time
+import argparse
+from io import BytesIO
+
+from contextlib import nullcontext
+import torch.distributed as dist
+
+from omegaconf import DictConfig, OmegaConf
+from torch.cuda.amp import autocast, GradScaler
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.algorithms.join import Join
+from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler
+from tensorboardX import SummaryWriter
+from funasr.train_utils.average_nbest_models import average_checkpoints
+
+from funasr.register import tables
+from funasr.optimizers import optim_classes
+from funasr.train_utils.trainer import Trainer
+from funasr.schedulers import scheduler_classes
+from funasr.train_utils.initialize import initialize
+from funasr.download.download_model_from_hub import download_model
+from funasr.models.lora.utils import mark_only_lora_as_trainable
+from funasr.train_utils.set_all_random_seed import set_all_random_seed
+from funasr.train_utils.load_pretrained_model import load_pretrained_model
+from funasr.utils.misc import prepare_model_dir
+from funasr.train_utils.model_summary import model_summary
+from funasr import AutoModel
+
+
+@hydra.main(config_name=None, version_base=None)
+def main_hydra(kwargs: DictConfig):
+    """Main hydra.
+    
+        Args:
+            kwargs: Additional keyword arguments.
+        """
+    if kwargs.get("debug", False):
+        import pdb
+
+        pdb.set_trace()
+
+    assert "model" in kwargs
+    if "model_conf" not in kwargs:
+        logging.info("download models from model hub: {}".format(kwargs.get("hub", "ms")))
+        kwargs = download_model(is_training=kwargs.get("is_training", True), **kwargs)
+
+    main(**kwargs)
+
+
+def main(**kwargs):
+
+    # set random seed
+    """Main.
+    
+        Args:
+            **kwargs: Additional keyword arguments.
+        """
+    set_all_random_seed(kwargs.get("seed", 0))
+    torch.backends.cudnn.enabled = kwargs.get("cudnn_enabled", torch.backends.cudnn.enabled)
+    torch.backends.cudnn.benchmark = kwargs.get("cudnn_benchmark", torch.backends.cudnn.benchmark)
+    torch.backends.cudnn.deterministic = kwargs.get("cudnn_deterministic", True)
+    # open tf32
+    torch.backends.cuda.matmul.allow_tf32 = kwargs.get("enable_tf32", True)
+
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    if local_rank == 0:
+        tables.print()
+    # Check if we are using DDP or FSDP
+    use_ddp = "WORLD_SIZE" in os.environ and int(os.environ["WORLD_SIZE"]) > 1
+    use_fsdp = kwargs.get("use_fsdp", False)
+    # use_ddp = False if use_fsdp else use_fsdp
+    if use_ddp or use_fsdp:
+        dist.init_process_group(backend=kwargs.get("backend", "nccl"), init_method="env://")
+        torch.cuda.set_device(local_rank)
+
+    logging.info("Build model, frontend, tokenizer")
+    device = kwargs.get("device", "cuda")
+    kwargs["device"] = "cpu"
+    model = AutoModel(**kwargs)
+
+    # save config.yaml
+    if (
+        (use_ddp or use_fsdp)
+        and dist.get_rank() == 0
+        or not (use_ddp or use_fsdp)
+        and local_rank == 0
+    ):
+        prepare_model_dir(**kwargs)
+
+    # parse kwargs
+    kwargs = model.kwargs
+    kwargs["device"] = device
+    tokenizer = kwargs["tokenizer"]
+    frontend = kwargs["frontend"]
+    model = model.model
+    del kwargs["model"]
+
+    # freeze_param
+    freeze_param = kwargs.get("freeze_param", None)
+    if freeze_param is not None:
+        if "," in freeze_param:
+            freeze_param = freeze_param.split(",")
+        if not isinstance(freeze_param, (list, tuple)):
+            freeze_param = (freeze_param,)
+        logging.info("freeze_param is not None: %s", freeze_param)
+        for t in freeze_param:
+            for k, p in model.named_parameters():
+                if k.startswith(t + ".") or k == t:
+                    logging.info(f"Setting {k}.requires_grad = False")
+                    p.requires_grad = False
+    lora_only = kwargs.get("lora_only", False)
+    if lora_only:
+        lora_bias = kwargs.get("lora_bias", "none")
+        logging.info("Enable LoRA-only training with bias=%s", lora_bias)
+        mark_only_lora_as_trainable(model, bias=lora_bias)
+    if local_rank == 0:
+        logging.info(f"{model_summary(model)}")
+
+    if use_ddp:
+        model = model.cuda(local_rank)
+        model = DDP(
+            model,
+            device_ids=[local_rank],
+            find_unused_parameters=kwargs.get("train_conf", {}).get(
+                "find_unused_parameters", False
+            ),
+        )
+    elif use_fsdp:
+        # model = FSDP(model).cuda(local_rank)
+
+        def custom_auto_wrap_policy(
+            module: nn.Module,
+            recurse: bool,
+            nonwrapped_numel: int,
+            # Additional custom arguments
+            min_num_params: int = int(1e8),
+        ) -> bool:
+            # 根据自定义逻辑决定是否包装模块
+            """Custom auto wrap policy.
+            
+                Args:
+                    module: TODO.
+                    recurse: TODO.
+                    nonwrapped_numel: TODO.
+                    min_num_params: TODO.
+                """
+            is_large = nonwrapped_numel >= min_num_params
+            requires_grad_uniform = len({p.requires_grad for p in module.parameters()}) == 1
+            return is_large and requires_grad_uniform
+
+        # Configure a custom `min_num_params`
+        my_auto_wrap_policy = functools.partial(custom_auto_wrap_policy, min_num_params=int(1e5))
+        torch.cuda.set_device(local_rank)
+        model = FSDP(
+            model,
+            auto_wrap_policy=custom_auto_wrap_policy,
+            mixed_precision=None,
+            device_id=torch.cuda.current_device(),
+        )
+    else:
+        model = model.to(device=kwargs.get("device", "cuda"))
+
+    kwargs["device"] = next(model.parameters()).device
+
+    # optim
+    logging.info("Build optim")
+    optim = kwargs.get("optim", "adam")
+    assert optim in optim_classes
+    optim_class = optim_classes.get(optim)
+    optim = optim_class(model.parameters(), **kwargs.get("optim_conf"))
+
+    # scheduler
+    logging.info("Build scheduler")
+    scheduler = kwargs.get("scheduler", "warmuplr")
+    assert scheduler in scheduler_classes
+    scheduler_class = scheduler_classes.get(scheduler)
+    scheduler = scheduler_class(optim, **kwargs.get("scheduler_conf"))
+
+    # dataset
+    logging.info("Build dataloader")
+    dataloader_class = tables.dataloader_classes.get(
+        kwargs["dataset_conf"].get("dataloader", "DataloaderMapStyle")
+    )
+    dataloader = dataloader_class(**kwargs)
+    # dataloader_tr, dataloader_val = dataloader_class(**kwargs)
+    trainer = Trainer(
+        local_rank=local_rank,
+        use_ddp=use_ddp,
+        use_fsdp=use_fsdp,
+        device=kwargs["device"],
+        output_dir=kwargs.get("output_dir", "./exp"),
+        **kwargs.get("train_conf"),
+    )
+
+    scaler = GradScaler(enabled=trainer.use_fp16) if trainer.use_fp16 else None
+    scaler = ShardedGradScaler(enabled=trainer.use_fp16) if trainer.use_fsdp else scaler
+
+    trainer.resume_checkpoint(
+        model=model,
+        optim=optim,
+        scheduler=scheduler,
+        scaler=scaler,
+    )
+
+    tensorboard_dir = os.path.join(kwargs.get("output_dir"), "tensorboard")
+    os.makedirs(tensorboard_dir, exist_ok=True)
+    try:
+        writer = SummaryWriter(tensorboard_dir)  # if trainer.rank == 0 else None
+    except:
+        writer = None
+
+    dataloader_tr, dataloader_val = None, None
+    for epoch in range(trainer.start_epoch, trainer.max_epoch):
+        time1 = time.perf_counter()
+
+        for data_split_i in range(trainer.start_data_split_i, dataloader.data_split_num):
+            time_slice_i = time.perf_counter()
+            dataloader_tr, dataloader_val = dataloader.build_iter(
+                epoch, data_split_i=data_split_i, start_step=trainer.start_step
+            )
+
+            trainer.train_epoch(
+                model=model,
+                optim=optim,
+                scheduler=scheduler,
+                scaler=scaler,
+                dataloader_train=dataloader_tr,
+                dataloader_val=dataloader_val,
+                epoch=epoch,
+                writer=writer,
+                data_split_i=data_split_i,
+                data_split_num=dataloader.data_split_num,
+                start_step=trainer.start_step,
+            )
+            trainer.start_step = 0
+
+            device = next(model.parameters()).device
+            if device.type == "cuda":
+                with torch.cuda.device(device):
+                    torch.cuda.empty_cache()
+
+            time_escaped = (time.perf_counter() - time_slice_i) / 3600.0
+            logging.info(
+                f"rank: {local_rank}, "
+                f"time_escaped_epoch: {time_escaped:.3f} hours, "
+                f"estimated to finish {dataloader.data_split_num} data_slices, remaining: {dataloader.data_split_num-data_split_i} slices, {(dataloader.data_split_num-data_split_i)*time_escaped:.3f} hours, "
+                f"epoch: {trainer.max_epoch - epoch} epochs, {((trainer.max_epoch - epoch - 1)*dataloader.data_split_num + dataloader.data_split_num-data_split_i)*time_escaped:.3f} hours\n"
+            )
+
+        trainer.start_data_split_i = 0
+        trainer.validate_epoch(
+            model=model, dataloader_val=dataloader_val, epoch=epoch + 1, writer=writer
+        )
+        scheduler.step()
+        trainer.step_in_epoch = 0
+        trainer.save_checkpoint(
+            epoch + 1, model=model, optim=optim, scheduler=scheduler, scaler=scaler
+        )
+
+        time2 = time.perf_counter()
+        time_escaped = (time2 - time1) / 3600.0
+        logging.info(
+            f"rank: {local_rank}, "
+            f"time_escaped_epoch: {time_escaped:.3f} hours, "
+            f"estimated to finish {trainer.max_epoch} "
+            f"epoch: {(trainer.max_epoch - epoch) * time_escaped:.3f} hours\n"
+        )
+        trainer.train_acc_avg = 0.0
+        trainer.train_loss_avg = 0.0
+
+    if trainer.rank == 0:
+        average_checkpoints(trainer.output_dir, trainer.avg_nbest_model)
+
+    trainer.close()
+
+
+if __name__ == "__main__":
+    main_hydra()

funasr/bin/train_ds.py

@@ -0,0 +1,259 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+
+import os
+import sys
+import torch
+import torch.nn as nn
+import hydra
+import logging
+import time
+import argparse
+from io import BytesIO
+
+from contextlib import nullcontext
+import torch.distributed as dist
+
+from omegaconf import DictConfig, OmegaConf
+from torch.cuda.amp import autocast, GradScaler
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.algorithms.join import Join
+from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler
+from funasr.train_utils.average_nbest_models import average_checkpoints
+
+from funasr.register import tables
+from funasr.optimizers import optim_classes
+from funasr.train_utils.trainer_ds import Trainer
+from funasr.schedulers import scheduler_classes
+from funasr.train_utils.initialize import initialize
+from funasr.download.download_model_from_hub import download_model
+from funasr.models.lora.utils import mark_only_lora_as_trainable
+from funasr.train_utils.set_all_random_seed import set_all_random_seed
+from funasr.train_utils.load_pretrained_model import load_pretrained_model
+from funasr.utils.misc import prepare_model_dir
+from funasr.train_utils.model_summary import model_summary
+from funasr import AutoModel
+
+try:
+    import deepspeed
+except:
+    deepspeed = None
+
+
+@hydra.main(config_name=None, version_base=None)
+def main_hydra(kwargs: DictConfig):
+    """Main hydra.
+    
+        Args:
+            kwargs: Additional keyword arguments.
+        """
+    if kwargs.get("debug", False):
+        import pdb
+
+        pdb.set_trace()
+
+    assert "model" in kwargs
+    if "model_conf" not in kwargs:
+        logging.info("download models from model hub: {}".format(kwargs.get("hub", "ms")))
+        kwargs = download_model(is_training=kwargs.get("is_training", True), **kwargs)
+
+    main(**kwargs)
+
+
+def main(**kwargs):
+
+    # set random seed
+    """Main.
+    
+        Args:
+            **kwargs: Additional keyword arguments.
+        """
+    set_all_random_seed(kwargs.get("seed", 0))
+    torch.backends.cudnn.enabled = kwargs.get("cudnn_enabled", torch.backends.cudnn.enabled)
+    torch.backends.cudnn.benchmark = kwargs.get("cudnn_benchmark", torch.backends.cudnn.benchmark)
+    torch.backends.cudnn.deterministic = kwargs.get("cudnn_deterministic", True)
+    # open tf32
+    torch.backends.cuda.matmul.allow_tf32 = kwargs.get("enable_tf32", True)
+
+    rank = int(os.environ.get("RANK", 0))
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+
+    if local_rank == 0:
+        tables.print()
+
+    use_ddp = world_size > 1
+    use_fsdp = kwargs.get("use_fsdp", False)
+    use_deepspeed = kwargs.get("use_deepspeed", False)
+    if use_deepspeed:
+        logging.info(f"use_deepspeed: {use_deepspeed}")
+        deepspeed.init_distributed(dist_backend=kwargs.get("backend", "nccl"))
+    elif use_ddp or use_fsdp:
+        logging.info(f"use_ddp: {use_ddp}, use_fsdp: {use_fsdp}")
+        dist.init_process_group(
+            backend=kwargs.get("backend", "nccl"),
+            init_method="env://",
+        )
+        torch.cuda.set_device(local_rank)
+
+    # rank = dist.get_rank()
+
+    logging.info("Build model, frontend, tokenizer")
+    device = kwargs.get("device", "cuda")
+    kwargs["device"] = "cpu"
+    model = AutoModel(**kwargs)
+
+    # save config.yaml
+    if rank == 0:
+        prepare_model_dir(**kwargs)
+
+    # parse kwargs
+    kwargs = model.kwargs
+    kwargs["device"] = device
+    tokenizer = kwargs["tokenizer"]
+    frontend = kwargs["frontend"]
+    model = model.model
+    del kwargs["model"]
+
+    # freeze_param
+    freeze_param = kwargs.get("freeze_param", None)
+    if freeze_param is not None:
+        if "," in freeze_param:
+            freeze_param = freeze_param.split(",")
+        if not isinstance(freeze_param, (list, tuple)):
+            freeze_param = (freeze_param,)
+        logging.info("freeze_param is not None: %s", freeze_param)
+        for t in freeze_param:
+            for k, p in model.named_parameters():
+                if k.startswith(t + ".") or k == t:
+                    logging.info(f"Setting {k}.requires_grad = False")
+                    p.requires_grad = False
+    lora_only = kwargs.get("lora_only", False)
+    if lora_only:
+        lora_bias = kwargs.get("lora_bias", "none")
+        logging.info("Enable LoRA-only training with bias=%s", lora_bias)
+        mark_only_lora_as_trainable(model, bias=lora_bias)
+    if local_rank == 0:
+        logging.info(f"{model_summary(model)}")
+
+    trainer = Trainer(
+        rank=rank,
+        local_rank=local_rank,
+        world_size=world_size,
+        use_ddp=use_ddp,
+        use_fsdp=use_fsdp,
+        device=kwargs["device"],
+        excludes=kwargs.get("excludes", None),
+        output_dir=kwargs.get("output_dir", "./exp"),
+        **kwargs.get("train_conf"),
+    )
+
+    model = trainer.warp_model(model, **kwargs)
+
+    kwargs["device"] = int(os.environ.get("LOCAL_RANK", 0))
+    trainer.device = int(os.environ.get("LOCAL_RANK", 0))
+
+    model, optim, scheduler = trainer.warp_optim_scheduler(model, **kwargs)
+
+    # dataset
+    logging.info("Build dataloader")
+    dataloader_class = tables.dataloader_classes.get(
+        kwargs["dataset_conf"].get("dataloader", "DataloaderMapStyle")
+    )
+    dataloader = dataloader_class(**kwargs)
+    # dataloader_tr, dataloader_val = dataloader_class(**kwargs)
+
+    scaler = GradScaler(enabled=True) if trainer.use_fp16 else None
+    scaler = ShardedGradScaler(enabled=trainer.use_fp16) if trainer.use_fsdp else scaler
+
+    trainer.resume_checkpoint(
+        model=model,
+        optim=optim,
+        scheduler=scheduler,
+        scaler=scaler,
+    )
+
+    early_stopping_patience = kwargs.get("train_conf", {}).get("early_stopping_patience", 0)
+    best_val_loss = float("inf")
+    epochs_no_improve = 0
+
+    dataloader_tr, dataloader_val = None, None
+    for epoch in range(trainer.start_epoch, trainer.max_epoch):
+        time1 = time.perf_counter()
+
+        for data_split_i in range(trainer.start_data_split_i, dataloader.data_split_num):
+            time_slice_i = time.perf_counter()
+
+            dataloader_tr, dataloader_val = dataloader.build_iter(
+                epoch, data_split_i=data_split_i, start_step=trainer.start_step
+            )
+
+            trainer.train_epoch(
+                model=model,
+                optim=optim,
+                scheduler=scheduler,
+                scaler=scaler,
+                dataloader_train=dataloader_tr,
+                dataloader_val=dataloader_val,
+                epoch=epoch,
+                data_split_i=data_split_i,
+                data_split_num=dataloader.data_split_num,
+                start_step=trainer.start_step,
+            )
+            trainer.start_step = 0
+
+            device = next(model.parameters()).device
+            if device.type == "cuda":
+                with torch.cuda.device(device):
+                    torch.cuda.empty_cache()
+
+            time_escaped = (time.perf_counter() - time_slice_i) / 3600.0
+            logging.info(
+                f"\n\nrank: {local_rank}, "
+                f"time_escaped_epoch: {time_escaped:.3f} hours, "
+                f"estimated to finish {dataloader.data_split_num} data_slices, remaining: {dataloader.data_split_num-data_split_i} slices, {(dataloader.data_split_num-data_split_i)*time_escaped:.3f} hours, "
+                f"epoch: {trainer.max_epoch - epoch} epochs, {((trainer.max_epoch - epoch - 1)*dataloader.data_split_num + dataloader.data_split_num-data_split_i)*time_escaped:.3f} hours\n"
+            )
+
+        trainer.start_data_split_i = 0
+        trainer.validate_epoch(model=model, dataloader_val=dataloader_val, epoch=epoch + 1)
+        current_val = trainer.val_loss_avg
+
+        if current_val < best_val_loss:
+            logging.info(f"current_val: {current_val}, best_val_loss: {best_val_loss}")
+            best_val_loss = current_val
+            epochs_no_improve = 0
+        else:
+            epochs_no_improve += 1
+            logging.info(f"No val_loss improvement for {epochs_no_improve}/{early_stopping_patience} epochs")
+        if early_stopping_patience > 0 and epochs_no_improve >= early_stopping_patience:
+            logging.info(f"Early stopping triggered at epoch {epoch+1}")
+            break
+
+        trainer.step_in_epoch = 0
+        trainer.save_checkpoint(
+            epoch + 1, model=model, optim=optim, scheduler=scheduler, scaler=scaler
+        )
+
+        time2 = time.perf_counter()
+        time_escaped = (time2 - time1) / 3600.0
+        logging.info(
+            f"\n\nrank: {local_rank}, "
+            f"time_escaped_epoch: {time_escaped:.3f} hours, "
+            f"estimated to finish {trainer.max_epoch} "
+            f"epoch: {(trainer.max_epoch - epoch) * time_escaped:.3f} hours\n"
+        )
+        trainer.train_acc_avg = 0.0
+        trainer.train_loss_avg = 0.0
+
+    if trainer.rank == 0:
+        average_checkpoints(
+            trainer.output_dir, trainer.avg_nbest_model, use_deepspeed=trainer.use_deepspeed
+        )
+
+    trainer.close()
+
+
+if __name__ == "__main__":
+    main_hydra()

funasr/cli.py

@@ -0,0 +1,174 @@
+"""FunASR CLI - Agent-friendly speech recognition from the command line."""
+
+import argparse
+import json
+import os
+import re
+import sys
+import time
+
+MODEL_CONFIGS = {
+    "sensevoice": {"model": "iic/SenseVoiceSmall", "vad_model": "fsmn-vad", "vad_kwargs": {"max_single_segment_time": 30000}},
+    "paraformer": {"model": "paraformer-zh", "vad_model": "fsmn-vad", "punc_model": "ct-punc"},
+    "paraformer-en": {"model": "paraformer-en", "vad_model": "fsmn-vad"},
+    "fun-asr-nano": {"model": "FunAudioLLM/Fun-ASR-Nano-2512", "vad_model": "fsmn-vad"},
+}
+
+
+def clean_text(text):
+    return re.sub(r"<\|[^|]*\|>", "", text).strip()
+
+
+def _srt_time(ms):
+    s = ms / 1000.0
+    h, m, sec = int(s // 3600), int((s % 3600) // 60), int(s % 60)
+    return f"{h:02d}:{m:02d}:{sec:02d},{int((s % 1) * 1000):03d}"
+
+
+def format_srt(segments):
+    lines = []
+    for i, seg in enumerate(segments, 1):
+        lines += [str(i), f"{_srt_time(seg.get('start',0))} --> {_srt_time(seg.get('end',0))}", seg.get('text',''), ""]
+    return "\n".join(lines)
+
+
+def format_tsv(segments):
+    lines = ["start\tend\ttext"]
+    for seg in segments:
+        lines.append(f"{seg.get('start',0)/1000:.3f}\t{seg.get('end',0)/1000:.3f}\t{seg.get('text','')}")
+    return "\n".join(lines)
+
+
+def _format_output(text, segments, timestamps, fmt, audio_path, model_name, language, elapsed):
+    if fmt == "text":
+        return text
+    elif fmt == "json":
+        obj = {"text": text}
+        if segments:
+            obj["segments"] = segments
+        if timestamps:
+            obj["timestamps"] = timestamps
+        try:
+            import soundfile as sf
+            audio_dur = round(sf.info(audio_path).duration, 3)
+        except Exception:
+            audio_dur = None
+        obj.update({"file": os.path.basename(audio_path), "model": model_name, "language": language or "auto", "audio_duration_s": audio_dur, "processing_s": round(elapsed, 3)})
+        return json.dumps(obj, ensure_ascii=False, indent=2)
+    elif fmt == "srt":
+        if segments:
+            return format_srt(segments)
+        # No per-sentence timestamps: emit one valid cue spanning the whole audio
+        # (instead of a bogus 99:59:59 end time).
+        try:
+            import soundfile as sf
+            dur_ms = int(sf.info(audio_path).duration * 1000)
+        except Exception:
+            dur_ms = 0
+        return f"1\n00:00:00,000 --> {_srt_time(dur_ms)}\n{text}\n"
+    elif fmt == "tsv":
+        return format_tsv(segments) if segments else f"start\tend\ttext\n0.000\t0.000\t{text}"
+
+
+def _get_version():
+    try:
+        from funasr import __version__
+        return __version__
+    except Exception:
+        return "unknown"
+
+
+def main():
+    p = argparse.ArgumentParser(
+        prog="funasr",
+        description="FunASR - speech recognition CLI. 50+ languages, speaker diarization.",
+        epilog="Examples:\n"
+               "  funasr audio.wav\n"
+               "  funasr audio.wav --model sensevoice -f json\n"
+               "  funasr audio.wav -f srt -o ./subs\n"
+               "  funasr audio.wav --spk --timestamps\n"
+               "  funasr audio.wav --hub hf --model fun-asr-nano\n",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    p.add_argument("audio", nargs="+", help="Audio file(s) to transcribe")
+    p.add_argument("--model", "-m", default="sensevoice", choices=list(MODEL_CONFIGS), help="Model (default: sensevoice)")
+    p.add_argument("--hub", "-H", default="ms", choices=["ms", "hf"], help="Model hub: ms (ModelScope) or hf (Hugging Face). Default: ms")
+    p.add_argument("--language", "-l", default=None, help="Language: zh, en, ja, ko, yue, auto")
+    p.add_argument("--device", default=None, help="Device: cuda:0, cpu (default: auto)")
+    p.add_argument("--output-format", "-f", default="text", choices=["text", "json", "srt", "tsv"], help="Output format (default: text)")
+    p.add_argument("--output-dir", "-o", default=None, help="Write output files to directory")
+    p.add_argument("--timestamps", action="store_true", help="Include word-level timestamps")
+    p.add_argument("--spk", action="store_true", help="Enable speaker diarization")
+    p.add_argument("--hotwords", default=None, help="Comma-separated hotwords")
+    p.add_argument("--verbose", "-v", action="store_true", help="Show loading/timing info on stderr")
+    p.add_argument("--version", action="version", version=f"%(prog)s {_get_version()}")
+    args = p.parse_args()
+
+    if args.verbose:
+        print(f"Loading model: {args.model} ...", file=sys.stderr)
+
+    import torch
+    from funasr import AutoModel
+
+    device = args.device or ("cuda:0" if torch.cuda.is_available() else "cpu")
+    config = MODEL_CONFIGS[args.model].copy()
+    config["hub"] = args.hub
+    if args.spk and "spk_model" not in config:
+        config["spk_model"] = "cam++"
+    if "punc_model" not in config and args.model not in ("fun-asr-nano", "sensevoice"):
+        config["punc_model"] = "ct-punc"
+
+    t_load = time.time()
+    model = AutoModel(device=device, disable_update=True, **config)
+    if args.verbose:
+        print(f"Model loaded in {time.time() - t_load:.1f}s", file=sys.stderr)
+
+    if args.output_dir:
+        os.makedirs(args.output_dir, exist_ok=True)
+
+    for audio_path in args.audio:
+        if not os.path.isfile(audio_path):
+            print(f"Error: file not found: {audio_path}", file=sys.stderr)
+            sys.exit(1)
+
+        if args.verbose:
+            print(f"Transcribing: {audio_path}", file=sys.stderr)
+
+        t0 = time.time()
+        gen_kw = {"input": audio_path, "batch_size": 1}
+        if args.language:
+            gen_kw["language"] = args.language
+        if args.hotwords:
+            gen_kw["hotwords"] = args.hotwords.split(",")
+
+        result = model.generate(**gen_kw)
+        elapsed = time.time() - t0
+
+        text = clean_text(result[0].get("text", ""))
+        segments = []
+        if "sentence_info" in result[0]:
+            for seg in result[0]["sentence_info"]:
+                s = {"start": seg.get("start", 0), "end": seg.get("end", 0), "text": clean_text(seg.get("sentence") or seg.get("text", ""))}
+                if args.spk and "spk" in seg:
+                    s["speaker"] = seg["spk"]
+                segments.append(s)
+
+        timestamps = result[0].get("timestamps") if args.timestamps else None
+        output = _format_output(text, segments, timestamps, args.output_format, audio_path, args.model, args.language, elapsed)
+
+        if args.output_dir:
+            ext = {"text": "txt", "json": "json", "srt": "srt", "tsv": "tsv"}[args.output_format]
+            out_path = os.path.join(args.output_dir, os.path.splitext(os.path.basename(audio_path))[0] + "." + ext)
+            with open(out_path, "w", encoding="utf-8") as f:
+                f.write(output)
+            if args.verbose:
+                print(f"Written: {out_path}", file=sys.stderr)
+        else:
+            print(output)
+
+        if args.verbose:
+            print(f"Done in {elapsed:.2f}s", file=sys.stderr)
+
+
+if __name__ == "__main__":
+    main()

funasr/datasets/audio_datasets/datasets.py

@@ -0,0 +1,330 @@
+import torch
+import random
+
+
+from funasr.register import tables
+from funasr.utils.load_utils import extract_fbank, load_audio_text_image_video
+
+
+@tables.register("dataset_classes", "AudioDataset")
+class AudioDataset(torch.utils.data.Dataset):
+    """
+    AudioDataset
+    """
+
+    def __init__(
+        self,
+        path,
+        index_ds: str = None,
+        frontend=None,
+        tokenizer=None,
+        is_training: bool = True,
+        int_pad_value: int = -1,
+        float_pad_value: float = 0.0,
+        **kwargs,
+    ):
+        """Initialize AudioDataset.
+        
+            Args:
+                path: TODO.
+                index_ds: TODO.
+                frontend: Audio frontend for feature extraction.
+                tokenizer: Tokenizer instance for text encoding/decoding.
+                is_training: Boolean flag for training.
+                int_pad_value: TODO.
+                float_pad_value: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        index_ds_class = tables.index_ds_classes.get(index_ds)
+        self.index_ds = index_ds_class(path, **kwargs)
+
+        self.preprocessor_speech = None
+        self.preprocessor_text = None
+
+        if is_training:
+            preprocessor_speech = kwargs.get("preprocessor_speech", None)
+            if preprocessor_speech:
+                preprocessor_speech_class = tables.preprocessor_classes.get(preprocessor_speech)
+                preprocessor_speech = preprocessor_speech_class(
+                    **kwargs.get("preprocessor_speech_conf")
+                )
+            self.preprocessor_speech = preprocessor_speech
+            preprocessor_text = kwargs.get("preprocessor_text", None)
+            if preprocessor_text:
+                preprocessor_text_class = tables.preprocessor_classes.get(preprocessor_text)
+                preprocessor_text = preprocessor_text_class(**kwargs.get("preprocessor_text_conf"))
+            self.preprocessor_text = preprocessor_text
+
+        self.frontend = frontend
+        self.fs = 16000 if frontend is None else frontend.fs
+        self.data_type = "sound"
+        self.tokenizer = tokenizer
+
+        self.int_pad_value = int_pad_value
+        self.float_pad_value = float_pad_value
+
+    def get_source_len(self, index):
+        """Get source len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_source_len(item)
+
+    def get_target_len(self, index):
+        """Get target len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_target_len(item)
+
+    def __len__(self):
+        """Internal: len  ."""
+        return len(self.index_ds)
+
+    def __getitem__(self, index):
+        """Internal: getitem  .
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        # import pdb;
+        # pdb.set_trace()
+        source = item["source"]
+        data_src = load_audio_text_image_video(source, fs=self.fs)
+        if self.preprocessor_speech:
+            data_src = self.preprocessor_speech(data_src, fs=self.fs)
+
+        speech, speech_lengths = extract_fbank(
+            data_src, data_type=self.data_type, frontend=self.frontend, is_final=True
+        )  # speech: [b, T, d]
+
+        target = item["target"]
+        if self.preprocessor_text:
+            target = self.preprocessor_text(target)
+
+        if self.tokenizer:
+            ids = self.tokenizer.encode(target)
+            text = torch.tensor(ids, dtype=torch.int64)
+        else:
+            ids = target
+            text = ids
+        ids_lengths = len(ids)
+        text_lengths = torch.tensor([ids_lengths], dtype=torch.int32)
+
+        return {
+            "speech": speech[0, :, :],
+            "speech_lengths": speech_lengths,
+            "text": text,
+            "text_lengths": text_lengths,
+        }
+
+    def collator(self, samples: list = None):
+        """Collator.
+        
+            Args:
+                samples: TODO.
+            """
+        outputs = {}
+        for sample in samples:
+            for key in sample.keys():
+                if key not in outputs:
+                    outputs[key] = []
+                outputs[key].append(sample[key])
+
+        for key, data_list in outputs.items():
+            if isinstance(data_list[0], torch.Tensor):
+                if data_list[0].dtype == torch.int64 or data_list[0].dtype == torch.int32:
+
+                    pad_value = self.int_pad_value
+                else:
+                    pad_value = self.float_pad_value
+
+                outputs[key] = torch.nn.utils.rnn.pad_sequence(
+                    data_list, batch_first=True, padding_value=pad_value
+                )
+        return outputs
+
+
+@tables.register("dataset_classes", "AudioDatasetHotword")
+class AudioDatasetHotword(AudioDataset):
+    # for finetuning contextual_paraformer and seaco_paraformer
+    def __init__(
+        self,
+        *args,
+        seaco_id: bool = 0,
+        **kwargs,
+    ):
+        """Initialize AudioDatasetHotword.
+        
+            Args:
+                *args: Variable positional arguments.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__(*args, **kwargs)
+        self.seaco_id = seaco_id
+
+    def __getitem__(self, index):
+        """Internal: getitem  .
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        # import pdb;
+        # pdb.set_trace()
+        source = item["source"]
+        data_src = load_audio_text_image_video(source, fs=self.fs)
+        if self.preprocessor_speech:
+            data_src = self.preprocessor_speech(data_src, fs=self.fs)
+        speech, speech_lengths = extract_fbank(
+            data_src, data_type=self.data_type, frontend=self.frontend, is_final=True
+        )  # speech: [b, T, d]
+
+        target = item["target"]
+        if self.preprocessor_text:
+            target = self.preprocessor_text(target)
+        if self.tokenizer:
+            ids = self.tokenizer.encode(target)
+            text = torch.tensor(ids, dtype=torch.int64)
+        else:
+            ids = target
+            text = ids
+        ids_lengths = len(ids)
+        text_lengths = torch.tensor([ids_lengths], dtype=torch.int32)
+
+        def generate_index(
+            length,
+            hotword_min_length=2,
+            hotword_max_length=8,
+            sample_rate=0.75,
+            double_rate=0.1,
+            pre_prob=0.0,
+            pre_index=None,
+            pre_hwlist=None,
+        ):
+            """Generate index.
+            
+                Args:
+                    length: TODO.
+                    hotword_min_length: TODO.
+                    hotword_max_length: TODO.
+                    sample_rate: TODO.
+                    double_rate: TODO.
+                    pre_prob: TODO.
+                    pre_index: TODO.
+                    pre_hwlist: TODO.
+                """
+            if length < hotword_min_length:
+                return [-1]
+            if random.random() < sample_rate:
+                if pre_prob > 0 and random.random() < pre_prob and pre_index is not None:
+                    return pre_index
+                if length == hotword_min_length:
+                    return [0, length - 1]
+                elif (
+                    random.random() < double_rate
+                    and length > hotword_max_length + hotword_min_length + 2
+                ):
+                    # sample two hotwords in a sentence
+                    _max_hw_length = min(hotword_max_length, length // 2)
+                    # first hotword
+                    start1 = random.randint(0, length // 3)
+                    end1 = random.randint(
+                        start1 + hotword_min_length - 1, start1 + _max_hw_length - 1
+                    )
+                    # second hotword
+                    start2 = random.randint(end1 + 1, length - hotword_min_length)
+                    end2 = random.randint(
+                        min(length - 1, start2 + hotword_min_length - 1),
+                        min(length - 1, start2 + hotword_max_length - 1),
+                    )
+                    return [start1, end1, start2, end2]
+                else:  # single hotword
+                    start = random.randint(0, length - hotword_min_length)
+                    end = random.randint(
+                        min(length - 1, start + hotword_min_length - 1),
+                        min(length - 1, start + hotword_max_length - 1),
+                    )
+                    return [start, end]
+            else:
+                return [-1]
+
+        hotword_indx = generate_index(text_lengths[0])
+        return {
+            "speech": speech[0, :, :],
+            "speech_lengths": speech_lengths,
+            "text": text,
+            "text_lengths": text_lengths,
+            "hotword_indx": hotword_indx,
+            "seaco_id": self.seaco_id,
+        }
+
+    def collator(self, samples: list = None):
+        """Collator.
+        
+            Args:
+                samples: TODO.
+            """
+        outputs = {}
+        hotword_indxs = []
+        seaco_id = samples[0]["seaco_id"]
+        for sample in samples:
+            for key in sample.keys():
+                if key == "seaco_id":
+                    continue
+                elif key == "hotword_indx":
+                    hotword_indxs.append(sample[key])
+                else:
+                    if key not in outputs:
+                        outputs[key] = []
+                    outputs[key].append(sample[key])
+
+        for key, data_list in outputs.items():
+            if isinstance(data_list[0], torch.Tensor):
+                if data_list[0].dtype == torch.int64 or data_list[0].dtype == torch.int32:
+                    pad_value = self.int_pad_value
+                else:
+                    pad_value = self.float_pad_value
+                outputs[key] = torch.nn.utils.rnn.pad_sequence(
+                    data_list, batch_first=True, padding_value=pad_value
+                )
+
+        hotword_list, hotword_lengths = [], []
+        text = outputs["text"]
+        seaco_label_pad = torch.ones_like(text) * -1 if seaco_id else None
+        for b, (hotword_indx, one_text, length) in enumerate(
+            zip(hotword_indxs, text, outputs["text_lengths"])
+        ):
+            length = length[0]
+            if seaco_label_pad is not None:
+                seaco_label_pad[b][:length] = seaco_id
+            if hotword_indx[0] != -1:
+                start, end = int(hotword_indx[0]), int(hotword_indx[1])
+                hotword = one_text[start : end + 1]
+                hotword_list.append(hotword)
+                hotword_lengths.append(end - start + 1)
+                if seaco_label_pad is not None:
+                    seaco_label_pad[b][start : end + 1] = one_text[start : end + 1]
+                if len(hotword_indx) == 4 and hotword_indx[2] != -1:
+                    # the second hotword if exist
+                    start, end = int(hotword_indx[2]), int(hotword_indx[3])
+                    hotword_list.append(one_text[start : end + 1])
+                    hotword_lengths.append(end - start + 1)
+                    if seaco_label_pad is not None:
+                        seaco_label_pad[b][start : end + 1] = one_text[start : end + 1]
+        hotword_list.append(torch.tensor([1]))
+        hotword_lengths.append(1)
+        hotword_pad = torch.nn.utils.rnn.pad_sequence(
+            hotword_list, batch_first=True, padding_value=0
+        )
+        outputs["hotword_pad"] = hotword_pad
+        outputs["hotword_lengths"] = torch.tensor(hotword_lengths, dtype=torch.int32)
+        if seaco_label_pad is not None:
+            outputs["seaco_label_pad"] = seaco_label_pad
+        return outputs

funasr/datasets/audio_datasets/espnet_samplers.py

@@ -0,0 +1,198 @@
+import torch
+import numpy as np
+import logging
+import math
+import torch.distributed as dist
+from torch.utils.data import DistributedSampler
+from torch.utils.data import BatchSampler, Sampler
+import torch.distributed as dist
+import random
+from funasr.register import tables
+
+
+@tables.register("batch_sampler_classes", "EspnetStyleBatchSampler")
+def EspnetStyleBatchSampler_fn(dataset, **kwargs):
+    """Espnetstylebatchsampler fn.
+    
+        Args:
+            dataset: TODO.
+            **kwargs: Additional keyword arguments.
+        """
+    dataloader_args = {}
+
+    batch_sampler = EspnetStyleBatchSampler(dataset, **kwargs)
+    dataloader_args["batch_sampler"] = batch_sampler
+    dataloader_args["num_workers"] = kwargs.get("num_workers", 4)
+    dataloader_args["pin_memory"] = kwargs.get("pin_memory", True)
+    num_workers = dataloader_args.get("num_workers", 4)
+    if num_workers > 0:
+        dataloader_args["persistent_workers"] = kwargs.get("persistent_workers", True)
+        dataloader_args["prefetch_factor"] = kwargs.get("prefetch_factor", 2)
+
+    return dataloader_args
+
+
+import torch
+from torch.utils.data import Dataset, DistributedSampler
+import math
+import random
+
+
+class EspnetStyleBatchSampler(DistributedSampler):
+    def __init__(
+        self,
+        dataset,
+        batch_size,
+        batch_type="token",
+        rank=None,
+        num_replicas=None,
+        rank_split=False,
+        shuffle=True,
+        drop_last=False,
+        is_training: bool = True,
+        sort_size: int = 1024,
+        start_step: int = 0,
+        **kwargs,
+    ):
+
+        """Initialize EspnetStyleBatchSampler.
+        
+            Args:
+                dataset: TODO.
+                batch_size: Number of samples per batch.
+                batch_type: TODO.
+                rank: TODO.
+                num_replicas: TODO.
+                rank_split: TODO.
+                shuffle: TODO.
+                drop_last: TODO.
+                is_training: Boolean flag for training.
+                sort_size: Size/dimension parameter.
+                start_step: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        try:
+            rank = dist.get_rank()
+            num_replicas = dist.get_world_size()
+        except:
+            rank = 0
+            num_replicas = 1
+        # if rank_split:
+        #     logging.info(f"Warning, rank_split: {rank_split}, batch and shuffle data in local rank")
+        #     rank = 0
+        #     num_replicas = 1
+        self.rank = rank
+        self.num_replicas = num_replicas
+        self.dataset = dataset
+        self.batch_size = batch_size
+        self.batch_type = batch_type
+        self.is_training = is_training
+        self.shuffle = shuffle and is_training
+        self.drop_last = drop_last
+
+        self.total_size = len(self.dataset)
+        self.num_samples = int(math.ceil(self.total_size / self.num_replicas))
+        self.epoch = 0
+        self.sort_size = sort_size * num_replicas
+        self.max_token_length = kwargs.get("max_token_length", 2048)
+        self.min_token_length = kwargs.get("min_token_length", 0)
+        self.length_scale_source = kwargs.get("length_scale_source", 1.0)
+        self.start_step = start_step
+        self.batch_num = 1
+        if self.start_step > 0:
+            logging.info(f"Warning, start_step > 0, dataloader start from step: {self.start_step}")
+        # super().__init__(dataset, num_replicas=num_replicas, rank=rank,
+        #                  shuffle=shuffle, drop_last=drop_last)
+
+    def __iter__(self):
+        """Internal: iter  ."""
+        if self.shuffle:
+            g = torch.Generator()
+            g.manual_seed(self.epoch)
+            random.seed(self.epoch)
+            indices = torch.randperm(len(self.dataset), generator=g).tolist()
+        else:
+            indices = list(range(len(self.dataset)))
+
+        # Sort indices by sample length
+        sorted_indices = sorted(indices, key=lambda idx: self.dataset.get_source_len(idx))
+
+        # Organize batches based on 'length' or 'example'
+        buffer_batches = []
+        batch = []
+        max_len_in_batch = 0  # Tracks the max sample length within the current batch
+
+        for idx in sorted_indices:
+
+            # original_sample_length = self.dataset.get_source_len(idx)
+            # if (
+            #     original_sample_length < self.min_token_length
+            #     or original_sample_length > self.max_token_length
+            # ):  # Skip samples that exceed the max length
+            #     continue
+
+            # sample_length = 1 if self.batch_type == "example" else original_sample_length
+
+            # Set sample_length based on the batch type
+            if self.batch_type == "example":
+                sample_length = 1
+            elif self.batch_type == "token":
+                sample_length = self.dataset.get_source_len(idx) + int(
+                    self.dataset.get_target_len(idx) * 1.2
+                )
+            else:
+                sample_length = self.dataset.get_source_len(idx)
+            # Calculate potential batch size with the new sample
+            potential_batch_length = max(max_len_in_batch, sample_length) * (len(batch) + 1)
+            # Add index to batch if it doesn't exceed batch size limit
+            if potential_batch_length <= self.batch_size:
+                batch.append(idx)
+                max_len_in_batch = max(max_len_in_batch, sample_length)
+            else:
+                # Save the current batch and start a new one
+                buffer_batches.append(batch)
+                batch = [idx]
+                max_len_in_batch = sample_length
+
+        # Add the last batch if it shouldn't be dropped
+        if batch and (not self.drop_last or len(batch) * max_len_in_batch == self.batch_size):
+            buffer_batches.append(batch)
+
+        # Shuffle the list of batches
+        if self.shuffle:
+            random.seed(self.epoch)
+            random.shuffle(buffer_batches)
+
+        # Ensure each rank gets the same number of batches
+        batches_per_rank = int(math.ceil(len(buffer_batches) / self.num_replicas))
+        total_batches_needed = batches_per_rank * self.num_replicas
+        extra_batches = total_batches_needed - len(buffer_batches)
+        # Add extra batches by random selection, if needed
+        buffer_batches += random.choices(buffer_batches, k=extra_batches)
+
+        # Allocate the batches to the current rank
+        start_idx = self.rank * batches_per_rank
+        end_idx = start_idx + batches_per_rank
+        rank_batches = buffer_batches[start_idx + self.start_step : end_idx]
+
+        self.batch_num = len(rank_batches)
+
+        logging.info(
+            f"rank: {self.rank}, dataloader start from step: {self.start_step}, batch_num: {end_idx-start_idx}, batch_num_after_step: {len(rank_batches)}"
+        )
+        # Return an iterator over the batches for the current rank
+        return iter(rank_batches)
+
+    def __len__(self):
+        # Calculate the number of batches per epoch for the current rank
+        """Internal: len  ."""
+        return self.batch_num
+
+    def set_epoch(self, epoch):
+        # Set the epoch for shuffling
+        """Set epoch.
+        
+            Args:
+                epoch: TODO.
+            """
+        self.epoch = epoch

funasr/datasets/audio_datasets/index_ds.py

@@ -0,0 +1,173 @@
+import os
+import json
+import torch
+import logging
+
+import librosa
+import random
+import torch.distributed as dist
+
+from funasr.register import tables
+
+
+@tables.register("index_ds_classes", "IndexDSJsonl")
+@tables.register("index_ds_classes", "IndexDSJsonlRankFull")
+@tables.register("index_ds_classes", "IndexDSJsonlRankSplit")
+class IndexDSJsonlRankFull(torch.utils.data.Dataset):
+
+    def __init__(self, path: str, **kwargs):
+        """Initialize IndexDSJsonlRankFull.
+        
+            Args:
+                path: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        self.max_source_length = kwargs.get("max_source_length", 2048)
+        self.min_source_length = kwargs.get("min_source_length", 0)
+        self.max_target_length = kwargs.get("max_target_length", 2048)
+        self.min_target_length = kwargs.get("min_target_length", 0)
+        self.max_token_length = kwargs.get("max_token_length", 2200)
+
+        is_training = kwargs.get("is_training", True)
+        if not (path.endswith(".jsonl") or path.endswith(".json")):
+            # jsonl list file
+            data_split_num = kwargs.get("data_split_num", 1)
+            data_split_i = kwargs.get("data_split_i", 0)
+
+            if not is_training:
+                data_split_num = 1
+                data_split_i = 0
+            with open(path, encoding="utf-8") as fin:
+                file_list_all = fin.readlines()
+
+                num_per_slice = (len(file_list_all) - 1) // data_split_num + 1  # 16
+                file_list = file_list_all[
+                    data_split_i * num_per_slice : (data_split_i + 1) * num_per_slice
+                ]
+                logging.info(
+                    f"is_training: {is_training}, data_split_num: {data_split_num}, data_split_i: {data_split_i}, \nfile_list: {file_list}, \nfile_list_all: {file_list_all}"
+                )
+
+        else:
+            file_list = [path]
+
+        # total_num = len(file_list)
+        # try:
+        #     rank = dist.get_rank()
+        #     world_size = dist.get_world_size()
+        # except:
+        #     rank = 0
+        #     world_size = 1
+        #     logging.info("distributed is not initialized, only single shard")
+        #
+        # if not kwargs.get("rank_split", False):
+        #     logging.info(f"Warning, rank_split disenabled, batch and shuffle data in global")
+        #     rank = 0
+        #     world_size = 1
+        #
+        # num_per_rank = total_num // world_size
+        # if num_per_rank * world_size < total_num:
+        #     logging.info(f"Warning, jsonl file:{total_num} could not be divided by world_size: {world_size}, {path}")
+        #     total_num_needed = num_per_rank * world_size
+        #
+        #     extra_num = total_num_needed - total_num
+        #     file_list_tmp = random.choices(file_list, k=extra_num)
+        #     file_list += file_list_tmp
+        #     logging.info(f"Warning, after random choices: {file_list}")
+        #
+        # file_list_rank = file_list[rank * num_per_rank:(rank + 1) * num_per_rank]
+        #
+        # logging.info(
+        #     f"is_training: {is_training}, file_list_rank: {file_list_rank}")
+
+        # contents = []
+        # for file_json in file_list_rank:
+        contents = []
+        for file_json in file_list:
+            with open(file_json.strip(), encoding="utf-8") as fin:
+                for line in fin:
+                    data = json.loads(line.strip())
+                    if "text" in data:  # for sft
+                        contents.append(data["text"])
+                    if "source" in data:  # for speech lab pretrain
+                        prompt = data.get("prompt", "<ASR>")
+                        source = data["source"].replace(
+                            "/cpfs01", "/cpfs_speech/data"
+                        )  # only use in alibaba gpu group: .replace("/cpfs01", "/cpfs_speech/data")
+                        target = data["target"]
+                        source_len = data.get("source_len", 1)
+                        target_len = data.get("target_len", 0)
+                        text_language = data.get("text_language", "")
+                        if "aishell" in source and text_language != "en":
+                            target = target.replace(" ", "")
+                        if (
+                            source_len < self.min_source_length
+                            or source_len > self.max_source_length
+                        ):
+                            continue
+                        if (
+                            target_len < self.min_target_length
+                            or target_len > self.max_target_length
+                        ):
+                            continue
+
+                        if (source_len + target_len) > self.max_token_length:
+                            continue
+
+                        contents_i = {
+                            "source": source,
+                            "prompt": prompt,
+                            "target": target,
+                            "source_len": source_len,
+                            "target_len": target_len,
+                        }
+                        text_language = data.get("text_language", None)
+                        if text_language is not None:
+                            contents_i["text_language"] = text_language
+                        if "emo_target" in data:
+                            contents_i["emo_target"] = data["emo_target"]
+                        if "event_target" in data:
+                            contents_i["event_target"] = data["event_target"]
+                        if "with_or_wo_itn" in data:
+                            contents_i["with_or_wo_itn"] = data["with_or_wo_itn"]
+                        # audio_language = data.get("audio_language", None)
+                        # if audio_language is not None:
+                        #     contents_i["audio_language"] = audio_language
+                        contents.append(contents_i)
+
+        self.contents = contents
+
+        logging.info("total_num of samplers: {}, {}".format(len(self.contents), path))
+
+    def __len__(self):
+        """Internal: len  ."""
+        return len(self.contents)
+
+    def __getitem__(self, index):
+
+        """Internal: getitem  .
+        
+            Args:
+                index: TODO.
+            """
+        data = self.contents[index]
+
+        return data
+
+    def get_source_len(self, data_dict):
+        """Get source len.
+        
+            Args:
+                data_dict: TODO.
+            """
+        return data_dict.get("source_len", 1)
+
+    def get_target_len(self, data_dict):
+
+        """Get target len.
+        
+            Args:
+                data_dict: TODO.
+            """
+        return data_dict.get("target_len", 0)

funasr/datasets/audio_datasets/jsonl2scp.py

@@ -0,0 +1,77 @@
+import os
+import json
+import torch
+import logging
+import hydra
+from omegaconf import DictConfig, OmegaConf
+import concurrent.futures
+import librosa
+import torch.distributed as dist
+
+
+def gen_scp_from_jsonl(jsonl_file, data_type_list, wav_scp_file, text_file):
+
+    """Gen scp from jsonl.
+    
+        Args:
+            jsonl_file: TODO.
+            data_type_list: TODO.
+            wav_scp_file: TODO.
+            text_file: TODO.
+        """
+    wav_f = open(wav_scp_file, "w")
+    text_f = open(text_file, "w")
+    with open(jsonl_file, encoding="utf-8") as fin:
+        for line in fin:
+            data = json.loads(line.strip())
+
+            prompt = data.get("prompt", "<ASR>")
+            source = data[data_type_list[0]]
+            target = data[data_type_list[1]]
+            source_len = data.get("source_len", 1)
+            target_len = data.get("target_len", 0)
+            if "aishell" in source:
+                target = target.replace(" ", "")
+            key = data["key"]
+            wav_f.write(f"{key}\t{source}\n")
+            wav_f.flush()
+            text_f.write(f"{key}\t{target}\n")
+            text_f.flush()
+
+    wav_f.close()
+    text_f.close()
+
+
+@hydra.main(config_name=None, version_base=None)
+def main_hydra(cfg: DictConfig):
+
+    """Main hydra.
+    
+        Args:
+            cfg: Configuration overrides.
+        """
+    kwargs = OmegaConf.to_container(cfg, resolve=True)
+    print(kwargs)
+
+    scp_file_list = kwargs.get(
+        "scp_file_list",
+        ("/Users/zhifu/funasr1.0/test_local/wav.scp", "/Users/zhifu/funasr1.0/test_local/text.txt"),
+    )
+    if isinstance(scp_file_list, str):
+        scp_file_list = eval(scp_file_list)
+    data_type_list = kwargs.get("data_type_list", ("source", "target"))
+    jsonl_file = kwargs.get(
+        "jsonl_file_in", "/Users/zhifu/funasr1.0/test_local/audio_datasets.jsonl"
+    )
+    gen_scp_from_jsonl(jsonl_file, data_type_list, *scp_file_list)
+
+
+"""
+python -m funasr.datasets.audio_datasets.json2scp \
+++scp_file_list='["/Users/zhifu/funasr1.0/test_local/wav.scp", "/Users/zhifu/funasr1.0/test_local/text.txt"]' \
+++data_type_list='["source", "target"]' \
+++jsonl_file_in=/Users/zhifu/funasr1.0/test_local/audio_datasets.jsonl
+"""
+
+if __name__ == "__main__":
+    main_hydra()

funasr/datasets/audio_datasets/preprocessor.py

@@ -0,0 +1,82 @@
+import os
+import json
+import torch
+import logging
+import concurrent.futures
+import librosa
+import torch.distributed as dist
+from typing import Collection
+import torch
+import torchaudio
+from torch import nn
+import random
+import re
+from funasr.tokenizer.cleaner import TextCleaner
+from funasr.register import tables
+
+
+@tables.register("preprocessor_classes", "SpeechPreprocessSpeedPerturb")
+class SpeechPreprocessSpeedPerturb(nn.Module):
+    def __init__(self, speed_perturb: list = None, **kwargs):
+        """Initialize SpeechPreprocessSpeedPerturb.
+        
+            Args:
+                speed_perturb: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        self.speed_perturb = speed_perturb
+
+    def forward(self, waveform, fs, **kwargs):
+        """Forward pass for training.
+        
+            Args:
+                waveform: TODO.
+                fs: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        if self.speed_perturb is None:
+            return waveform
+        speed = random.choice(self.speed_perturb)
+        if speed != 1.0:
+            if not isinstance(waveform, torch.Tensor):
+                waveform = torch.tensor(waveform)
+            waveform, _ = torchaudio.sox_effects.apply_effects_tensor(
+                waveform.view(1, -1), fs, [["speed", str(speed)], ["rate", str(fs)]]
+            )
+            waveform = waveform.view(-1)
+
+        return waveform
+
+
+@tables.register("preprocessor_classes", "TextPreprocessSegDict")
+class TextPreprocessSegDict(nn.Module):
+    def __init__(
+        self,
+        seg_dict: str = None,
+        text_cleaner: Collection[str] = None,
+        split_with_space: bool = False,
+        **kwargs
+    ):
+        """Initialize TextPreprocessSegDict.
+        
+            Args:
+                seg_dict: TODO.
+                text_cleaner: TODO.
+                split_with_space: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+
+        self.text_cleaner = TextCleaner(text_cleaner)
+
+    def forward(self, text, **kwargs):
+        """Forward pass for training.
+        
+            Args:
+                text: Text tensor or string input.
+                **kwargs: Additional keyword arguments.
+            """
+        text = self.text_cleaner(text)
+
+        return text

funasr/datasets/audio_datasets/samplers.py

@@ -0,0 +1,592 @@
+import torch
+import numpy as np
+import logging
+import math
+import random
+import torch.distributed as dist
+from torch.utils.data import DistributedSampler
+from torch.utils.data import BatchSampler, Sampler
+import torch.distributed as dist
+
+from funasr.register import tables
+
+
+@tables.register("batch_sampler_classes", "BatchSampler")
+@tables.register("batch_sampler_classes", "CustomDistributedBatchSampler")
+@tables.register("batch_sampler_classes", "CustomDistributedDynamicBatchSampler")
+@tables.register("batch_sampler_classes", "DynamicBatchLocalShuffleSampler")
+@tables.register("batch_sampler_classes", "RankFullLocalShuffleBatchSampler")
+@tables.register("batch_sampler_classes", "RankFullLocalShuffleDynamicBatchSampler")
+def CustomDistributedBatchSampler_fn(dataset, **kwargs):
+    """Customdistributedbatchsampler fn.
+    
+        Args:
+            dataset: TODO.
+            **kwargs: Additional keyword arguments.
+        """
+    dataloader_args = {}
+    batch_type = kwargs.get("batch_type", "example")
+    if batch_type == "example":
+        batch_sampler = CustomDistributedBatchSampler(dataset, **kwargs)
+
+    else:
+        if kwargs.get("sort_size", -1) > 0:
+            batch_sampler = CustomDistributedBufferDynamicBatchSampler(dataset, **kwargs)
+        else:
+            batch_sampler = CustomDistributedDynamicBatchSampler(dataset, **kwargs)
+        # batch_sampler = CustomDistributedDynamicBatchSampler(dataset, **kwargs)
+
+    dataloader_args["batch_sampler"] = batch_sampler
+    dataloader_args["num_workers"] = kwargs.get("num_workers", 4)
+    dataloader_args["pin_memory"] = kwargs.get("pin_memory", True)
+    num_workers = dataloader_args.get("num_workers", 4)
+    if num_workers > 0:
+        dataloader_args["persistent_workers"] = kwargs.get("persistent_workers", True)
+        dataloader_args["prefetch_factor"] = kwargs.get("prefetch_factor", 2)
+
+    return dataloader_args
+
+
+class CustomDistributedBatchSampler(Sampler):
+    def __init__(
+        self,
+        dataset,
+        batch_size,
+        num_replicas=None,
+        rank=None,
+        shuffle=True,
+        drop_last=False,
+        is_training: bool = True,
+        **kwargs,
+    ):
+
+        """Initialize CustomDistributedBatchSampler.
+        
+            Args:
+                dataset: TODO.
+                batch_size: Number of samples per batch.
+                num_replicas: TODO.
+                rank: TODO.
+                shuffle: TODO.
+                drop_last: TODO.
+                is_training: Boolean flag for training.
+                **kwargs: Additional keyword arguments.
+            """
+        try:
+            rank = dist.get_rank()
+            num_replicas = dist.get_world_size()
+        except:
+            rank = 0
+            num_replicas = 1
+        self.rank = rank
+        self.num_replicas = num_replicas
+        self.dataset = dataset
+        self.batch_size = batch_size
+        self.is_training = is_training
+        self.shuffle = shuffle and is_training
+        self.drop_last = drop_last
+        # self.total_size = len(dataset)
+        if self.drop_last:
+            self.total_size = (len(self.dataset) // (batch_size * num_replicas)) * (
+                batch_size * num_replicas
+            )
+        else:
+            self.total_size = math.ceil(len(self.dataset) / (batch_size * num_replicas)) * (
+                batch_size * num_replicas
+            )
+        self.num_samples = int(self.total_size // self.num_replicas)
+        self.epoch = 0
+        self.max_token_length = kwargs.get("max_token_length", None)
+        self.length_scale_source = kwargs.get("length_scale_source", 1.0)
+
+    def __iter__(self):
+        # Generate a list of indices
+        """Internal: iter  ."""
+        if self.shuffle:
+            g = torch.Generator()
+            g.manual_seed(self.epoch)
+            indices = torch.randperm(len(self.dataset), generator=g).tolist()
+        else:
+            indices = list(range(len(self.dataset)))
+
+        # Add extra samples to make it evenly divisible
+        padding_size = self.total_size - len(indices)
+        if padding_size <= len(indices):
+            indices += indices[:padding_size]
+        else:
+            indices += (
+                indices * (padding_size // len(indices)) + indices[: padding_size % len(indices)]
+            )
+
+        assert len(indices) == self.total_size
+
+        # Subsample
+        indices = indices[self.rank : self.total_size : self.num_replicas]
+        assert len(indices) == self.num_samples
+
+        # Filter out indices with length greater than the max length, if provided
+        if self.max_token_length is not None:
+            filtered_indices = []
+            for idx in indices:
+                source_len = self.dataset.get_source_len(idx) / self.length_scale_source
+                if source_len <= self.max_token_length:
+                    filtered_indices.append(idx)
+            indices = filtered_indices
+
+        # Now that we have only the indices for this replica, chunk them into batches
+        batches = [
+            indices[i : i + self.batch_size] for i in range(0, len(indices), self.batch_size)
+        ]
+
+        # Drop the last batch if it's not full and drop_last is True
+        if self.drop_last and len(batches[-1]) != self.batch_size:
+            batches = batches[:-1]
+
+        return iter(batches)
+
+    def __len__(self):
+
+        """Internal: len  ."""
+        return self.num_samples // self.batch_size
+
+    def set_epoch(self, epoch):
+        """Set epoch.
+        
+            Args:
+                epoch: TODO.
+            """
+        self.epoch = epoch
+
+
+class CustomDistributedBufferBatchSampler(Sampler):
+    def __init__(
+        self,
+        dataset,
+        batch_size,
+        num_replicas=None,
+        rank=None,
+        shuffle=True,
+        drop_last=False,
+        is_training: bool = True,
+        sort_size: int = 1024,
+        **kwargs,
+    ):
+
+        """Initialize CustomDistributedBufferBatchSampler.
+        
+            Args:
+                dataset: TODO.
+                batch_size: Number of samples per batch.
+                num_replicas: TODO.
+                rank: TODO.
+                shuffle: TODO.
+                drop_last: TODO.
+                is_training: Boolean flag for training.
+                sort_size: Size/dimension parameter.
+                **kwargs: Additional keyword arguments.
+            """
+        try:
+            rank = dist.get_rank()
+            num_replicas = dist.get_world_size()
+        except:
+            rank = 0
+            num_replicas = 1
+        self.rank = rank
+        self.num_replicas = num_replicas
+        self.dataset = dataset
+        self.batch_size = batch_size
+        self.is_training = is_training
+        self.shuffle = shuffle and is_training
+        self.drop_last = drop_last
+        # self.total_size = len(dataset)
+        if self.drop_last:
+            self.total_size = (len(self.dataset) // (batch_size * num_replicas)) * (
+                batch_size * num_replicas
+            )
+        else:
+            self.total_size = math.ceil(len(self.dataset) / (batch_size * num_replicas)) * (
+                batch_size * num_replicas
+            )
+        self.num_samples = int(self.total_size // self.num_replicas)
+        self.epoch = 0
+        self.max_token_length = kwargs.get("max_token_length", None)
+        self.length_scale_source = kwargs.get("length_scale_source", 1.0)
+        self.sort_size = sort_size
+
+    def __iter__(self):
+        # Generate a list of indices
+        """Internal: iter  ."""
+        if self.shuffle:
+            g = torch.Generator()
+            g.manual_seed(self.epoch)
+            indices = torch.randperm(len(self.dataset), generator=g).tolist()
+        else:
+            indices = list(range(len(self.dataset)))
+
+        # Add extra samples to make it evenly divisible
+        padding_size = self.total_size - len(indices)
+        if padding_size <= len(indices):
+            indices += indices[:padding_size]
+        else:
+            indices += (
+                indices * (padding_size // len(indices)) + indices[: padding_size % len(indices)]
+            )
+
+        assert len(indices) == self.total_size
+
+        # Subsample
+        indices = indices[self.rank : self.total_size : self.num_replicas]
+        assert len(indices) == self.num_samples
+
+        # Filter out indices with length greater than the max length, if provided
+        if self.max_token_length is not None:
+            filtered_indices = []
+            for idx in indices:
+                source_len = self.dataset.get_source_len(idx) / self.length_scale_source
+                if source_len <= self.max_token_length:
+                    filtered_indices.append(idx)
+            indices = filtered_indices
+
+        # Buffer sorting logic
+        sorted_batches = []
+        buffer = []
+
+        for idx in indices:
+            buffer.append(idx)
+            if len(buffer) >= self.sort_size:
+                # Sort the buffer based on some criteria, e.g., dataset sample length
+                buffer.sort(key=lambda x: self.dataset.get_source_len(x))
+                sorted_batches.extend(self._create_batches_from_buffer(buffer))
+                buffer = []
+
+        # Handle the remaining items in the buffer
+        if buffer:
+            buffer.sort(key=lambda x: self.dataset.get_source_len(x))
+            sorted_batches.extend(self._create_batches_from_buffer(buffer))
+
+        return iter(sorted_batches)
+
+    def _create_batches_from_buffer(self, buffer):
+        # Function to convert the sorted buffer into batches
+        """Internal: create batches from buffer.
+        
+            Args:
+                buffer: TODO.
+            """
+        batched_buffer = [
+            buffer[i : i + self.batch_size] for i in range(0, len(buffer), self.batch_size)
+        ]
+        if self.drop_last and len(batched_buffer[-1]) != self.batch_size:
+            batched_buffer = batched_buffer[:-1]
+        return batched_buffer
+
+    def __len__(self):
+
+        """Internal: len  ."""
+        return self.num_samples // self.batch_size
+
+    def set_epoch(self, epoch):
+        """Set epoch.
+        
+            Args:
+                epoch: TODO.
+            """
+        self.epoch = epoch
+
+
+class CustomDistributedDynamicBatchSampler(DistributedSampler):
+    def __init__(
+        self,
+        dataset,
+        batch_size,
+        num_replicas=None,
+        rank=None,
+        shuffle=True,
+        drop_last=False,
+        is_training: bool = True,
+        **kwargs,
+    ):
+
+        """Initialize CustomDistributedDynamicBatchSampler.
+        
+            Args:
+                dataset: TODO.
+                batch_size: Number of samples per batch.
+                num_replicas: TODO.
+                rank: TODO.
+                shuffle: TODO.
+                drop_last: TODO.
+                is_training: Boolean flag for training.
+                **kwargs: Additional keyword arguments.
+            """
+        try:
+            rank = dist.get_rank()
+            num_replicas = dist.get_world_size()
+        except:
+            rank = 0
+            num_replicas = 1
+        self.rank = rank
+        self.num_replicas = num_replicas
+        self.dataset = dataset
+        self.batch_size = batch_size
+        self.is_training = is_training
+        self.shuffle = shuffle and is_training
+        self.drop_last = drop_last
+
+        self.total_size = len(self.dataset)
+        # self.num_samples = int(math.ceil(self.total_size / self.num_replicas))
+        self.epoch = 0
+        self.max_token_length = kwargs.get("max_token_length", 2048)
+        self.length_scale_source = kwargs.get("length_scale_source", 1.0)
+
+    def __iter__(self):
+        """Internal: iter  ."""
+        if self.shuffle:
+            g = torch.Generator()
+            g.manual_seed(self.epoch)
+            indices = torch.randperm(len(self.dataset), generator=g).tolist()
+        else:
+            indices = list(range(len(self.dataset)))
+
+        indices = indices[self.rank : self.total_size : self.num_replicas]
+
+        batches = []
+        batch = []
+        max_len_in_batch = 0
+        current_batch_length = 0
+
+        for idx in indices:
+            sample_length = self.dataset.get_source_len(idx)
+            if sample_length > self.max_token_length:
+                continue
+            potential_batch_length = (
+                max_len_in_batch if sample_length < max_len_in_batch else sample_length
+            ) * (len(batch) + 1)
+
+            if potential_batch_length <= self.batch_size:
+                batch.append(idx)
+                if sample_length > max_len_in_batch:
+                    max_len_in_batch = sample_length
+                    # current_batch_length = max_len_in_batch * len(batch)
+            else:
+                batches.append(batch)
+                batch = [idx]
+                max_len_in_batch = sample_length
+                # current_batch_length = max_len_in_batch
+
+        # Add the last batch if it's not empty and we're not dropping it
+        if batch and (not self.drop_last or len(batch) * max_len_in_batch == self.batch_size):
+            batches.append(batch)
+
+        return iter(batches)
+
+    def __len__(self):
+
+        """Internal: len  ."""
+        return 1
+
+    def set_epoch(self, epoch):
+        """Set epoch.
+        
+            Args:
+                epoch: TODO.
+            """
+        self.epoch = epoch
+
+
+class CustomDistributedBufferDynamicBatchSampler(DistributedSampler):
+    def __init__(
+        self,
+        dataset,
+        batch_size,
+        batch_type="token",
+        num_replicas=None,
+        rank=None,
+        rank_split=False,
+        shuffle=True,
+        drop_last=False,
+        is_training: bool = True,
+        sort_size: int = 1024,
+        start_step: int = 0,
+        **kwargs,
+    ):
+
+        """Initialize CustomDistributedBufferDynamicBatchSampler.
+        
+            Args:
+                dataset: TODO.
+                batch_size: Number of samples per batch.
+                batch_type: TODO.
+                num_replicas: TODO.
+                rank: TODO.
+                rank_split: TODO.
+                shuffle: TODO.
+                drop_last: TODO.
+                is_training: Boolean flag for training.
+                sort_size: Size/dimension parameter.
+                start_step: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        try:
+            rank = dist.get_rank()
+            num_replicas = dist.get_world_size()
+        except:
+            rank = 0
+            num_replicas = 1
+
+        # if rank_split:
+        #     logging.info(f"Warning, rank_split: {rank_split}, batch and shuffle data in local rank")
+        #     rank = 0
+        #     num_replicas = 1
+
+        self.rank = rank
+        self.num_replicas = num_replicas
+        self.dataset = dataset
+        self.batch_size = batch_size
+        self.batch_type = batch_type
+        self.is_training = is_training
+        self.shuffle = shuffle and is_training
+        self.drop_last = drop_last
+
+        self.total_size = len(self.dataset)
+        self.num_samples = int(math.ceil(self.total_size / self.num_replicas))
+        self.epoch = 0
+        self.sort_size = sort_size * num_replicas
+        self.max_token_length = kwargs.get("max_token_length", 2048)
+        self.length_scale_source = kwargs.get("length_scale_source", 1.0)
+        self.batch_size_sample_max = kwargs.get("batch_size_sample_max", 200)
+        self.start_step = start_step
+        self.batch_num = 1
+        if self.start_step > 0:
+            logging.info(f"Warning, start_step > 0, dataloader start from step: {self.start_step}")
+        # super().__init__(
+        #     dataset, num_replicas=num_replicas, rank=rank, shuffle=shuffle, drop_last=drop_last
+        # )
+
+    def __iter__(self):
+        """Internal: iter  ."""
+        if self.shuffle:
+            g = torch.Generator()
+            g.manual_seed(self.epoch)
+            random.seed(self.epoch)
+
+            indices = torch.randperm(len(self.dataset), generator=g).tolist()
+        else:
+            indices = list(range(len(self.dataset)))
+
+        # Create sorted buffers and form batches
+        buffer_batches = []
+        for i in range(0, len(indices), self.sort_size):
+            buffer = sorted(
+                indices[i : i + self.sort_size], key=lambda idx: self.dataset.get_source_len(idx)
+            )
+            batch = []
+            max_len_in_batch = 0
+            count = 1
+            for idx in buffer:
+                original_sample_length = self.dataset.get_source_len(idx)
+                if original_sample_length > self.max_token_length:
+                    continue
+                sample_length = 1 if self.batch_type == "example" else original_sample_length
+                potential_batch_length = max(max_len_in_batch, sample_length) * (len(batch) + 1)
+                if potential_batch_length <= self.batch_size and count < self.batch_size_sample_max:
+                    batch.append(idx)
+                    max_len_in_batch = max(max_len_in_batch, sample_length)
+                    count += 1
+                else:
+                    buffer_batches.append(batch)
+                    batch = [idx]
+                    max_len_in_batch = sample_length
+                    count = 1
+            if batch:
+                buffer_batches.append(batch)
+
+        # Ensure each rank gets the same number of batches, duplicate data if needed
+        batches_per_rank = math.ceil(len(buffer_batches) / self.num_replicas)
+        total_batches_needed = batches_per_rank * self.num_replicas
+
+        extra_batches = total_batches_needed - len(buffer_batches)
+        buffer_batches += random.choices(buffer_batches, k=extra_batches)
+
+        # Evenly distribute batches from buffer_batches to each rank
+        rank_batches = [[] for _ in range(self.num_replicas)]
+        for i, batch in enumerate(buffer_batches):
+            rank_batches[i % self.num_replicas].append(batch)
+
+        # Assign all batches for the current rank directly
+        final_batches = rank_batches[self.rank][self.start_step :]
+        self.batch_num = len(final_batches)
+
+        logging.info(
+            f"rank: {self.rank}, dataloader start from step: {self.start_step}, batch_num: {len(rank_batches[self.rank])}, after: {self.batch_num}"
+        )
+        return iter(final_batches)
+
+    def __len__(self):
+        # Calculate the number of batches per epoch for the current rank
+        """Internal: len  ."""
+        return self.batch_num
+
+    def set_epoch(self, epoch):
+        """Set epoch.
+        
+            Args:
+                epoch: TODO.
+            """
+        self.epoch = epoch
+
+
+class DistributedSamplerWarp(BatchSampler):
+    def __init__(
+        self, dataset, batch_size, num_replicas=None, rank=None, shuffle=True, drop_last=False
+    ):
+        """Initialize DistributedSamplerWarp.
+        
+            Args:
+                dataset: TODO.
+                batch_size: Number of samples per batch.
+                num_replicas: TODO.
+                rank: TODO.
+                shuffle: TODO.
+                drop_last: TODO.
+            """
+        if num_replicas is None:
+            if not torch.distributed.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            num_replicas = torch.distributed.get_world_size()
+        if rank is None:
+            if not torch.distributed.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            rank = torch.distributed.get_rank()
+
+        self.dataset = dataset
+        self.batch_size = batch_size
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.shuffle = shuffle
+        self.drop_last = drop_last
+
+        # Create an instance of the DistributedSampler
+        self.sampler = DistributedSampler(
+            self.dataset, num_replicas=self.num_replicas, rank=self.rank, shuffle=self.shuffle
+        )
+
+        # Call BatchSampler's constructor
+        super().__init__(self.sampler, batch_size, drop_last)
+
+    def __iter__(self):
+        # If we shuffle, we need to call the set_epoch method
+        """Internal: iter  ."""
+        if self.shuffle:
+            self.sampler.set_epoch(self.epoch)
+
+        # Generate batch indices using the parent class
+        return super().__iter__()
+
+    def set_epoch(self, epoch):
+        """Set epoch.
+        
+            Args:
+                epoch: TODO.
+            """
+        self.epoch = epoch

funasr/datasets/audio_datasets/scp2jsonl.py

@@ -0,0 +1,148 @@
+import os
+import json
+import torch
+import logging
+import hydra
+from omegaconf import DictConfig, OmegaConf
+import concurrent.futures
+import librosa
+import torch.distributed as dist
+from tqdm import tqdm
+
+
+def gen_jsonl_from_wav_text_list(
+    path, data_type_list=("source", "target"), jsonl_file_out: str = None, **kwargs
+):
+    """Gen jsonl from wav text list.
+    
+        Args:
+            path: TODO.
+            data_type_list: TODO.
+            jsonl_file_out: TODO.
+            **kwargs: Additional keyword arguments.
+        """
+    try:
+        rank = dist.get_rank()
+        world_size = dist.get_world_size()
+    except:
+        rank = 0
+        world_size = 1
+
+    cpu_cores = os.cpu_count() or 1
+    print(f"convert wav.scp text to jsonl, ncpu: {cpu_cores}")
+    if rank == 0:
+        json_dict = {}
+        for data_type, data_file in zip(data_type_list, path):
+            json_dict[data_type] = {}
+            with open(data_file, "r") as f:
+
+                data_file_lists = f.readlines()
+                lines_for_each_th = (len(data_file_lists) - 1) // cpu_cores + 1
+                task_num = cpu_cores if len(data_file_lists) > cpu_cores else 1
+                # import pdb;pdb.set_trace()
+                if task_num > 1:
+                    with concurrent.futures.ThreadPoolExecutor(max_workers=cpu_cores) as executor:
+
+                        futures = [
+                            executor.submit(
+                                parse_context_length,
+                                data_file_lists[
+                                    i * lines_for_each_th : (i + 1) * lines_for_each_th
+                                ],
+                                data_type,
+                                i,
+                            )
+                            for i in range(task_num)
+                        ]
+
+                        for future in concurrent.futures.as_completed(futures):
+
+                            json_dict[data_type].update(future.result())
+                else:
+                    res = parse_context_length(data_file_lists, data_type)
+                    json_dict[data_type].update(res)
+
+        with open(jsonl_file_out, "w") as f:
+            for key in json_dict[data_type_list[0]].keys():
+                jsonl_line = {"key": key}
+                for data_file in data_type_list:
+                    if key in json_dict[data_file]:
+                        jsonl_line.update(json_dict[data_file][key])
+                jsonl_line = json.dumps(jsonl_line, ensure_ascii=False)
+                f.write(jsonl_line + "\n")
+                f.flush()
+        print(f"processed {len(json_dict[data_type_list[0]])} samples")
+
+    else:
+        pass
+
+    if world_size > 1:
+        dist.barrier()
+
+
+def parse_context_length(data_list: list, data_type: str, id=0):
+    """Parse context length.
+    
+        Args:
+            data_list: TODO.
+            data_type: TODO.
+            id: TODO.
+        """
+    pbar = tqdm(total=len(data_list), dynamic_ncols=True)
+    res = {}
+    for i, line in enumerate(data_list):
+        pbar.update(1)
+        pbar.set_description(f"cpu: {id}")
+        lines = line.strip().split(maxsplit=1)
+        key = lines[0]
+        line = lines[1] if len(lines) > 1 else ""
+        line = line.strip()
+        if data_type == "source":
+            if os.path.exists(line):
+                waveform, _ = librosa.load(line, sr=16000)
+                sample_num = len(waveform)
+                context_len = int(sample_num * 1000 / 16000 / 10)
+            else:
+                print("source file not found: {}".format(line))
+                continue
+        else:
+            context_len = len(line.split()) if " " in line else len(line)
+        res[key] = {data_type: line, f"{data_type}_len": context_len}
+    return res
+
+
+@hydra.main(config_name=None, version_base=None)
+def main_hydra(cfg: DictConfig):
+
+    """Main hydra.
+    
+        Args:
+            cfg: Configuration overrides.
+        """
+    kwargs = OmegaConf.to_container(cfg, resolve=True)
+    print(kwargs)
+
+    scp_file_list = kwargs.get(
+        "scp_file_list",
+        ("/Users/zhifu/funasr1.0/test_local/wav.scp", "/Users/zhifu/funasr1.0/test_local/text.txt"),
+    )
+    if isinstance(scp_file_list, str):
+        scp_file_list = eval(scp_file_list)
+    data_type_list = kwargs.get("data_type_list", ("source", "target"))
+    jsonl_file_out = kwargs.get(
+        "jsonl_file_out", "/Users/zhifu/funasr1.0/test_local/audio_datasets.jsonl"
+    )
+    gen_jsonl_from_wav_text_list(
+        scp_file_list, data_type_list=data_type_list, jsonl_file_out=jsonl_file_out
+    )
+
+
+"""
+python -m funasr.datasets.audio_datasets.scp2jsonl \
+++scp_file_list='["/Users/zhifu/funasr1.0/test_local/wav.scp", "/Users/zhifu/funasr1.0/test_local/text.txt"]' \
+++data_type_list='["source", "target"]' \
+++jsonl_file_out=/Users/zhifu/funasr1.0/test_local/audio_datasets.jsonl
+"""
+
+if __name__ == "__main__":
+    main_hydra()

funasr/datasets/audio_datasets/scp2len.py

@@ -0,0 +1,141 @@
+import os
+import json
+import torch
+import logging
+import hydra
+from omegaconf import DictConfig, OmegaConf
+import concurrent.futures
+import librosa
+import torch.distributed as dist
+from tqdm import tqdm
+
+
+def gen_jsonl_from_wav_text_list(
+    path, data_type_list=("source",), jsonl_file_out: str = None, **kwargs
+):
+    """Gen jsonl from wav text list.
+    
+        Args:
+            path: TODO.
+            data_type_list: TODO.
+            jsonl_file_out: TODO.
+            **kwargs: Additional keyword arguments.
+        """
+    try:
+        rank = dist.get_rank()
+        world_size = dist.get_world_size()
+    except:
+        rank = 0
+        world_size = 1
+
+    cpu_cores = os.cpu_count() or 1
+    print(f"convert wav.scp text to jsonl, ncpu: {cpu_cores}")
+    if rank == 0:
+        json_dict = {}
+        # for data_type, data_file in zip(data_type_list, path):
+        data_type = data_type_list[0]
+        data_file = path
+        json_dict[data_type] = {}
+        with open(data_file, "r") as f:
+
+            data_file_lists = f.readlines()
+            print("")
+            lines_for_each_th = (len(data_file_lists) - 1) // cpu_cores + 1
+            task_num = cpu_cores if len(data_file_lists) > cpu_cores else 1
+            # import pdb;pdb.set_trace()
+            if task_num > 1:
+                with concurrent.futures.ThreadPoolExecutor(max_workers=cpu_cores) as executor:
+
+                    futures = [
+                        executor.submit(
+                            parse_context_length,
+                            data_file_lists[i * lines_for_each_th : (i + 1) * lines_for_each_th],
+                            data_type,
+                            i,
+                        )
+                        for i in range(task_num)
+                    ]
+
+                    for future in concurrent.futures.as_completed(futures):
+
+                        json_dict[data_type].update(future.result())
+            else:
+                res = parse_context_length(data_file_lists, data_type)
+                json_dict[data_type].update(res)
+
+        with open(jsonl_file_out, "w") as f:
+            for key in json_dict[data_type_list[0]].keys():
+                jsonl_line = {"key": key}
+                for data_file in data_type_list:
+                    jsonl_line.update(json_dict[data_file][key])
+                # jsonl_line = json.dumps(jsonl_line, ensure_ascii=False)
+                source_len = jsonl_line["source_len"]
+                jsonl_line = f"{key} {source_len}"
+                f.write(jsonl_line + "\n")
+                f.flush()
+        print(f"processed {len(json_dict[data_type_list[0]])} samples")
+
+    else:
+        pass
+
+    if world_size > 1:
+        dist.barrier()
+
+
+def parse_context_length(data_list: list, data_type: str, id=0):
+    """Parse context length.
+    
+        Args:
+            data_list: TODO.
+            data_type: TODO.
+            id: TODO.
+        """
+    pbar = tqdm(total=len(data_list), dynamic_ncols=True)
+    res = {}
+    for i, line in enumerate(data_list):
+        pbar.update(1)
+        pbar.set_description(f"cpu: {id}")
+        lines = line.strip().split(maxsplit=1)
+        key = lines[0]
+        line = lines[1] if len(lines) > 1 else ""
+        line = line.strip()
+        if os.path.exists(line):
+            waveform, _ = librosa.load(line, sr=16000)
+            sample_num = len(waveform)
+            context_len = int(sample_num / 16000 * 1000 / 10)
+        else:
+            context_len = len(line.split()) if " " in line else len(line)
+        res[key] = {data_type: line, f"{data_type}_len": context_len}
+    return res
+
+
+@hydra.main(config_name=None, version_base=None)
+def main_hydra(cfg: DictConfig):
+
+    """Main hydra.
+    
+        Args:
+            cfg: Configuration overrides.
+        """
+    kwargs = OmegaConf.to_container(cfg, resolve=True)
+    print(kwargs)
+
+    scp_file_list = kwargs.get("scp_file_list", "/Users/zhifu/funasr1.0/data/list/train_wav.scp")
+    # if isinstance(scp_file_list, str):
+    #     scp_file_list = eval(scp_file_list)
+    data_type_list = kwargs.get("data_type_list", ("source",))
+    jsonl_file_out = kwargs.get("jsonl_file_out", "/Users/zhifu/funasr1.0/data/list/wav_len.txt")
+    gen_jsonl_from_wav_text_list(
+        scp_file_list, data_type_list=data_type_list, jsonl_file_out=jsonl_file_out
+    )
+
+
+"""
+python -m funasr.datasets.audio_datasets.scp2jsonl \
+++scp_file_list='["/Users/zhifu/funasr1.0/test_local/wav.scp", "/Users/zhifu/funasr1.0/test_local/text.txt"]' \
+++data_type_list='["source", "target"]' \
+++jsonl_file_out=/Users/zhifu/funasr1.0/test_local/audio_datasets.jsonl
+"""
+
+if __name__ == "__main__":
+    main_hydra()

funasr/datasets/audio_datasets/sensevoice2jsonl.py

@@ -0,0 +1,216 @@
+import os
+import json
+import torch
+import logging
+import hydra
+import re
+import string
+from omegaconf import DictConfig, OmegaConf
+import concurrent.futures
+import librosa
+import torch.distributed as dist
+from tqdm import tqdm
+
+
+def gen_jsonl_from_wav_text_list(
+    path, data_type_list=("source", "target"), jsonl_file_out: str = None, model_dir: str = "iic/SenseVoiceSmall", **kwargs
+):
+    """Gen jsonl from wav text list.
+    
+        Args:
+            path: TODO.
+            data_type_list: TODO.
+            jsonl_file_out: TODO.
+            model_dir: TODO.
+            **kwargs: Additional keyword arguments.
+        """
+    try:
+        rank = dist.get_rank()
+        world_size = dist.get_world_size()
+    except:
+        rank = 0
+        world_size = 1
+
+    cpu_cores = os.cpu_count() or 1
+    print(f"convert wav.scp text to jsonl, ncpu: {cpu_cores}")
+    if rank == 0:
+        json_dict = {}
+        for data_type, data_file in zip(data_type_list, path):
+            json_dict[data_type] = {}
+            with open(data_file, "r") as f:
+
+                data_file_lists = f.readlines()
+                lines_for_each_th = (len(data_file_lists) - 1) // cpu_cores + 1
+                task_num = cpu_cores if len(data_file_lists) > cpu_cores else 1
+                # import pdb;pdb.set_trace()
+                if task_num > 1:
+                    with concurrent.futures.ThreadPoolExecutor(max_workers=cpu_cores) as executor:
+
+                        futures = [
+                            executor.submit(
+                                parse_context_length,
+                                data_file_lists[
+                                    i * lines_for_each_th : (i + 1) * lines_for_each_th
+                                ],
+                                data_type,
+                                i,
+                            )
+                            for i in range(task_num)
+                        ]
+
+                        for future in concurrent.futures.as_completed(futures):
+
+                            json_dict[data_type].update(future.result())
+                else:
+                    res = parse_context_length(data_file_lists, data_type)
+                    json_dict[data_type].update(res)
+
+        if "text_language" not in data_type_list or "emo_target" not in data_type_list or "event_target" not in data_type_list:
+            from funasr import AutoModel
+
+            model = AutoModel(
+                model=model_dir,
+            )
+
+            rich_dict = {}
+            for key in json_dict["source"].keys():
+                input_wav = json_dict["source"][key]["source"]
+                res = model.generate(
+                input=input_wav,
+                cache={},
+                language="auto",  # "zn", "en", "yue", "ja", "ko", "nospeech"
+                use_itn=True,
+                )
+                text = res[0]["text"]
+                pattern = r"<\|[^|]+\|>"
+                matches = re.findall(pattern, text)
+                text_language, emo_target, event_target = matches[:3]
+                rich_dict[key] = [text_language, emo_target, event_target]
+
+
+            if "text_language" not in data_type_list:
+                data_type_list.append("text_language")
+                if "text_language" not in json_dict:
+                    json_dict["text_language"] = {}
+                for key in json_dict["source"].keys():
+                    json_dict["text_language"][key] = {}
+                    json_dict["text_language"][key]["text_language"] = rich_dict[key][0]
+
+            if "emo_target" not in data_type_list:
+                data_type_list.append("emo_target")
+                if "emo_target" not in json_dict:
+                    json_dict["emo_target"] = {}
+                for key in json_dict["source"].keys():
+                    json_dict["emo_target"][key] = {}
+                    json_dict["emo_target"][key]["emo_target"] = rich_dict[key][1]
+
+            if "event_target" not in data_type_list:
+                data_type_list.append("event_target")
+                if "event_target" not in json_dict:
+                    json_dict["event_target"] = {}
+                for key in json_dict["source"].keys():
+                    json_dict["event_target"][key] = {}
+                    json_dict["event_target"][key]["event_target"] = rich_dict[key][2]
+
+        with open(jsonl_file_out, "w") as f:
+            for key in json_dict[data_type_list[0]].keys():
+                jsonl_line = {"key": key}
+                for data_file in data_type_list:
+                    jsonl_line.update(json_dict[data_file][key])
+                jsonl_line = json.dumps(jsonl_line, ensure_ascii=False)
+                f.write(jsonl_line + "\n")
+                f.flush()
+        print(f"processed {len(json_dict[data_type_list[0]])} samples")
+
+    else:
+        pass
+
+    if world_size > 1:
+        dist.barrier()
+
+def contains_punctuation(s):
+    """Contains punctuation.
+    
+        Args:
+            s: TODO.
+        """
+    punctuations = (
+        string.punctuation +  
+        '，。、；：？！""''（）【】《》〈〉「」『』〔〕［］｛｝～·…—–'  
+    )
+    return any(char in punctuations for char in s)
+
+def parse_context_length(data_list: list, data_type: str, id=0):
+    """Parse context length.
+    
+        Args:
+            data_list: TODO.
+            data_type: TODO.
+            id: TODO.
+        """
+    pbar = tqdm(total=len(data_list), dynamic_ncols=True)
+    res = {}
+    for i, line in enumerate(data_list):
+        pbar.update(1)
+        pbar.set_description(f"cpu: {id}")
+        lines = line.strip().split(maxsplit=1)
+        key = lines[0]
+        line = lines[1] if len(lines) > 1 else ""
+        line = line.strip()
+        if os.path.exists(line):
+            waveform, _ = librosa.load(line, sr=16000)
+            sample_num = len(waveform)
+            context_len = int(sample_num / 16000 * 1000 / 10)
+        else:
+            context_len = len(line.split()) if " " in line else len(line)
+        if data_type == "source":
+            res[key] = {data_type: line, f"{data_type}_len": context_len}
+        elif data_type == "target":
+            punc = contains_punctuation(line)
+            if punc:
+                with_or_wo_itn = "<|withitn|>"
+            else:
+                with_or_wo_itn = "<|woitn|>"
+            res[key] = {data_type: line, f"{data_type}_len": context_len, "with_or_wo_itn": with_or_wo_itn}
+        else:
+            res[key] = {data_type: line}
+    return res
+
+
+@hydra.main(config_name=None, version_base=None)
+def main_hydra(cfg: DictConfig):
+
+    """Main hydra.
+    
+        Args:
+            cfg: Configuration overrides.
+        """
+    kwargs = OmegaConf.to_container(cfg, resolve=True)
+    print(kwargs)
+
+    scp_file_list = kwargs.get(
+        "scp_file_list",
+        ("/Users/zhifu/funasr1.0/test_local/wav.scp", "/Users/zhifu/funasr1.0/test_local/text.txt"),
+    )
+    if isinstance(scp_file_list, str):
+        scp_file_list = eval(scp_file_list)
+    data_type_list = kwargs.get("data_type_list", ("source", "target"))
+    jsonl_file_out = kwargs.get(
+        "jsonl_file_out", "/Users/zhifu/funasr1.0/test_local/audio_datasets.jsonl"
+    )
+    model_dir = kwargs.get("model_dir", "iic/SenseVoiceSmall")
+    gen_jsonl_from_wav_text_list(
+        scp_file_list, data_type_list=data_type_list, jsonl_file_out=jsonl_file_out, model_dir=model_dir
+    )
+
+
+"""
+python -m funasr.datasets.audio_datasets.sensevoice2jsonl \
+++scp_file_list='["/Users/zhifu/funasr1.0/test_local/wav.scp", "/Users/zhifu/funasr1.0/test_local/text.txt", "/Users/zhifu/funasr1.0/test_local/text_language.txt", "/Users/zhifu/funasr1.0/test_local/emo_target.txt", "/Users/zhifu/funasr1.0/test_local/event_target.txt"]' \
+++data_type_list='["source", "target", "text_language", "emo_target", "event_target"]' \
+++jsonl_file_out='/Users/zhifu/funasr1.0/test_local/audio_datasets.jsonl' \
+++model_dir='iic/SenseVoiceSmall'
+"""
+
+if __name__ == "__main__":
+    main_hydra()

funasr/datasets/audio_datasets/update_jsonl.py

@@ -0,0 +1,124 @@
+import os
+import json
+import torch
+import logging
+import hydra
+from omegaconf import DictConfig, OmegaConf
+import concurrent.futures
+import librosa
+import torch.distributed as dist
+import threading
+from tqdm import tqdm
+from concurrent.futures import ThreadPoolExecutor
+
+
+def gen_scp_from_jsonl(jsonl_file, jsonl_file_out, ncpu):
+    """Gen scp from jsonl.
+    
+        Args:
+            jsonl_file: TODO.
+            jsonl_file_out: TODO.
+            ncpu: TODO.
+        """
+    jsonl_file_out_f = open(jsonl_file_out, "w")
+    with open(jsonl_file, encoding="utf-8") as fin:
+        lines = fin.readlines()
+
+        num_total = len(lines)
+        if ncpu > 1:
+            # 使用ThreadPoolExecutor限制并发线程数
+            with ThreadPoolExecutor(max_workers=ncpu) as executor:
+                # 提交任务到线程池
+                futures = {executor.submit(update_data, lines, i) for i in tqdm(range(num_total))}
+
+                # 等待所有任务完成，这会阻塞直到所有提交的任务完成
+                for future in concurrent.futures.as_completed(futures):
+                    # 这里可以添加额外的逻辑来处理完成的任务，但在这个例子中我们只是等待
+                    pass
+        else:
+            for i in range(num_total):
+                update_data(lines, i)
+        logging.info("All audio durations have been processed.")
+
+        for line in lines:
+
+            jsonl_file_out_f.write(line + "\n")
+            jsonl_file_out_f.flush()
+
+    jsonl_file_out_f.close()
+
+
+def update_data(lines, i):
+    """Update data.
+    
+        Args:
+            lines: TODO.
+            i: TODO.
+        """
+    line = lines[i]
+    data = json.loads(line.strip())
+
+    wav_path = data["source"].replace("/cpfs01", "/cpfs_speech/data")
+    if os.path.exists(wav_path):
+        waveform, _ = librosa.load(wav_path, sr=16000)
+        sample_num = len(waveform)
+        source_len = int(sample_num / 16000 * 1000 / 10)
+        source_len_old = data["source_len"]
+        # if (source_len_old - source_len) > 100 or (source_len - source_len_old) > 100:
+        #     logging.info(f"old: {source_len_old}, new: {source_len}, wav: {wav_path}")
+        data["source_len"] = source_len
+        data["source"] = wav_path
+        jsonl_line = json.dumps(data, ensure_ascii=False)
+        lines[i] = jsonl_line
+
+
+def update_wav_len(jsonl_file_list_in, jsonl_file_out_dir, ncpu=1):
+
+    """Update wav len.
+    
+        Args:
+            jsonl_file_list_in: TODO.
+            jsonl_file_out_dir: TODO.
+            ncpu: TODO.
+        """
+    os.makedirs(jsonl_file_out_dir, exist_ok=True)
+    with open(jsonl_file_list_in, "r") as f:
+        data_file_lists = f.readlines()
+
+        for i, jsonl in enumerate(data_file_lists):
+            filename_with_extension = os.path.basename(jsonl.strip())
+            jsonl_file_out = os.path.join(jsonl_file_out_dir, filename_with_extension)
+            logging.info(f"{i}/{len(data_file_lists)}, jsonl: {jsonl}, {jsonl_file_out}")
+
+            gen_scp_from_jsonl(jsonl.strip(), jsonl_file_out, ncpu)
+
+
+@hydra.main(config_name=None, version_base=None)
+def main_hydra(cfg: DictConfig):
+
+    """Main hydra.
+    
+        Args:
+            cfg: Configuration overrides.
+        """
+    kwargs = OmegaConf.to_container(cfg, resolve=True)
+    logging.info(kwargs)
+
+    jsonl_file_list_in = kwargs.get(
+        "jsonl_file_list_in", "/Users/zhifu/funasr1.0/data/list/data_jsonl.list"
+    )
+    jsonl_file_out_dir = kwargs.get("jsonl_file_out_dir", "/Users/zhifu/funasr1.0/data_tmp")
+    ncpu = kwargs.get("ncpu", 1)
+    update_wav_len(jsonl_file_list_in, jsonl_file_out_dir, ncpu)
+    # gen_scp_from_jsonl(jsonl_file_list_in, jsonl_file_out_dir)
+
+
+"""
+python -m funasr.datasets.audio_datasets.json2scp \
+++scp_file_list='["/Users/zhifu/funasr1.0/test_local/wav.scp", "/Users/zhifu/funasr1.0/test_local/text.txt"]' \
+++data_type_list='["source", "target"]' \
+++jsonl_file_in=/Users/zhifu/funasr1.0/test_local/audio_datasets.jsonl
+"""
+
+if __name__ == "__main__":
+    main_hydra()

funasr/datasets/dataloader_entry.py

@@ -0,0 +1,168 @@
+import logging
+import torch
+
+from funasr.register import tables
+
+
+# @tables.register("dataloader_classes", "DataloaderMapStyle")
+def DataloaderMapStyle(frontend=None, tokenizer=None, **kwargs):
+    # dataset
+    """Dataloadermapstyle.
+    
+        Args:
+            frontend: Audio frontend for feature extraction.
+            tokenizer: Tokenizer instance for text encoding/decoding.
+            **kwargs: Additional keyword arguments.
+        """
+    logging.info("Build dataloader")
+    dataset_class = tables.dataset_classes.get(kwargs.get("dataset", "AudioDataset"))
+    dataset_tr = dataset_class(
+        kwargs.get("train_data_set_list"),
+        frontend=frontend,
+        tokenizer=tokenizer,
+        is_training=True,
+        **kwargs.get("dataset_conf"),
+    )
+    dataset_val = dataset_class(
+        kwargs.get("valid_data_set_list"),
+        frontend=frontend,
+        tokenizer=tokenizer,
+        is_training=False,
+        **kwargs.get("dataset_conf"),
+    )
+
+    # dataloader
+    batch_sampler = kwargs["dataset_conf"].get("batch_sampler", "BatchSampler")
+    batch_sampler_val = None
+    if batch_sampler is not None:
+        batch_sampler_class = tables.batch_sampler_classes.get(batch_sampler)
+        batch_sampler = batch_sampler_class(dataset_tr, **kwargs.get("dataset_conf"))
+        batch_sampler_val = batch_sampler_class(
+            dataset_val, is_training=False, **kwargs.get("dataset_conf")
+        )
+
+    dataloader_tr = torch.utils.data.DataLoader(
+        dataset_tr, collate_fn=dataset_tr.collator, **batch_sampler
+    )
+    dataloader_val = torch.utils.data.DataLoader(
+        dataset_val, collate_fn=dataset_val.collator, **batch_sampler_val
+    )
+
+    return dataloader_tr, dataloader_val
+
+
+@tables.register("dataloader_classes", "DataloaderMapStyle")
+class DataloaderMapStyle:
+    def __init__(self, frontend=None, tokenizer=None, **kwargs):
+        # dataset
+        """Initialize DataloaderMapStyle.
+        
+            Args:
+                frontend: Audio frontend for feature extraction.
+                tokenizer: Tokenizer instance for text encoding/decoding.
+                **kwargs: Additional keyword arguments.
+            """
+        logging.info("Build dataloader")
+
+        dataset_class = tables.dataset_classes.get(kwargs.get("dataset", "AudioDataset"))
+        dataset_tr = None
+        # split dataset
+        self.data_split_num = kwargs["dataset_conf"].get("data_split_num", 1)
+        if self.data_split_num == 1:
+            dataset_tr = dataset_class(
+                kwargs.get("train_data_set_list"),
+                frontend=frontend,
+                tokenizer=tokenizer,
+                is_training=True,
+                **kwargs.get("dataset_conf"),
+            )
+        dataset_val = dataset_class(
+            kwargs.get("valid_data_set_list"),
+            frontend=frontend,
+            tokenizer=tokenizer,
+            is_training=False,
+            **kwargs.get("dataset_conf"),
+        )
+
+        self.dataset_tr = dataset_tr
+        self.dataset_val = dataset_val
+        self.kwargs = kwargs
+
+        self.dataset_class = dataset_class
+        self.frontend = frontend
+        self.tokenizer = tokenizer
+        self.kwargs = kwargs
+
+    def build_iter(self, epoch=0, data_split_i=0, start_step=0, **kwargs):
+
+        # reload dataset slice
+        """Build iter.
+        
+            Args:
+                epoch: TODO.
+                data_split_i: TODO.
+                start_step: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        if self.data_split_num > 1:
+            del self.dataset_tr
+            self.dataset_tr = self.dataset_class(
+                self.kwargs.get("train_data_set_list"),
+                frontend=self.frontend,
+                tokenizer=self.tokenizer,
+                is_training=True,
+                **self.kwargs.get("dataset_conf"),
+                data_split_i=data_split_i,
+            )
+
+        # dataloader
+        batch_sampler = self.kwargs["dataset_conf"].get("batch_sampler", "BatchSampler")
+        batch_sampler_val = None
+        if batch_sampler is not None:
+            batch_sampler_class = tables.batch_sampler_classes.get(batch_sampler)
+            batch_sampler = batch_sampler_class(
+                self.dataset_tr, start_step=start_step, **self.kwargs.get("dataset_conf")
+            )
+            batch_sampler_val = batch_sampler_class(
+                self.dataset_val, is_training=False, **self.kwargs.get("dataset_conf")
+            )
+
+        batch_sampler["batch_sampler"].set_epoch(epoch)
+        batch_sampler_val["batch_sampler"].set_epoch(epoch)
+        dataloader_tr = torch.utils.data.DataLoader(
+            self.dataset_tr, collate_fn=self.dataset_tr.collator, **batch_sampler
+        )
+        dataloader_val = torch.utils.data.DataLoader(
+            self.dataset_val, collate_fn=self.dataset_val.collator, **batch_sampler_val
+        )
+
+        return dataloader_tr, dataloader_val
+
+
+@tables.register("dataloader_classes", "DataloaderIterable")
+def DataloaderIterable(frontend=None, tokenizer=None, **kwargs):
+    """Dataloaderiterable.
+    
+        Args:
+            frontend: Audio frontend for feature extraction.
+            tokenizer: Tokenizer instance for text encoding/decoding.
+            **kwargs: Additional keyword arguments.
+        """
+    logging.info("Build dataloader")
+    dataset_class = tables.dataset_classes.get(kwargs.get("dataset", "LargeDataset"))
+    dataset_tr = dataset_class(
+        kwargs.get("train_data_set_list"),
+        frontend=frontend,
+        tokenizer=tokenizer,
+        is_training=True,
+        **kwargs.get("dataset_conf"),
+    )
+    dataset_val = dataset_class(
+        kwargs.get("valid_data_set_list"),
+        frontend=frontend,
+        tokenizer=tokenizer,
+        is_training=False,
+        **kwargs.get("dataset_conf"),
+    )
+
+    return dataset_tr, dataset_val

funasr/datasets/fun_asr_datasets/datasets.py

@@ -0,0 +1,569 @@
+import json
+import logging
+
+import re
+import torch
+import random
+import traceback
+import numpy as np
+from funasr.register import tables
+from funasr.utils.load_utils import extract_fbank, load_audio_text_image_video
+
+
+@tables.register("dataset_classes", "FunASR")
+class FunASR(torch.utils.data.Dataset):
+    """
+    FunASR dataset
+    """
+
+    def __init__(
+            self,
+            path,
+            index_ds: str = None,
+            frontend=None,
+            tokenizer=None,
+            int_pad_value: int = -1,
+            float_pad_value: float = 0.0,
+            **kwargs,
+    ):
+        """Initialize FunASR.
+        
+            Args:
+                path: TODO.
+                index_ds: TODO.
+                frontend: Audio frontend for feature extraction.
+                tokenizer: Tokenizer instance for text encoding/decoding.
+                int_pad_value: TODO.
+                float_pad_value: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        index_ds_class = tables.index_ds_classes.get(index_ds)
+        self.index_ds = index_ds_class(path, **kwargs)
+        preprocessor_speech = kwargs.get("preprocessor_speech", None)
+        if preprocessor_speech:
+            preprocessor_speech_class = tables.preprocessor_classes.get(preprocessor_speech)
+            preprocessor_speech = preprocessor_speech_class(
+                **kwargs.get("preprocessor_speech_conf")
+            )
+        self.preprocessor_speech = preprocessor_speech
+
+        preprocessor_noise = kwargs.get("preprocessor_noise", None)
+        if preprocessor_noise:
+            preprocessor_noise_class = tables.preprocessor_classes.get(preprocessor_noise)
+            preprocessor_noise = preprocessor_noise_class(**kwargs.get("preprocessor_noise_conf"))
+        self.preprocessor_noise = preprocessor_noise
+
+        prompt_classes_text = kwargs.get("prompt_classes", None)
+        if prompt_classes_text is not None:
+            prompt_classes = tables.prompt_classes.get(prompt_classes_text)
+            prompt_classes = prompt_classes(**kwargs.get("prompt_conf"))
+        else:
+            prompt_classes = None
+        self.prompt_classes = prompt_classes
+
+        self.frontend = frontend
+        self.fs = 16000 if frontend is None else frontend.fs
+        self.data_type = "sound"
+        self.tokenizer = tokenizer
+
+        self.int_pad_value = int_pad_value
+        self.float_pad_value = float_pad_value
+        self.sos = kwargs.get("sos", "<|startoftranscript|>")
+        self.eos = kwargs.get("eos", "<|endoftext|>")
+        self.batch_size = kwargs.get("batch_size")
+        self.batch_type = kwargs.get("batch_type")
+        self.prompt_ids_len = 0
+        self.retry = kwargs.get("retry", 100)
+
+        self.pattern = re.compile(r"(<\|startofspeech\|>.*?<\|endofspeech\|>)")
+        # self.kwargs = kwargs
+        self.max_token_length = kwargs.get("max_token_length", 1500)
+        self.batch_size_scale_ratio_max = kwargs.get("batch_size_scale_ratio_max", 1.5)
+        self.batch_size_token_max = kwargs.get("batch_size_token_max", 2500)
+        self.multiturn_num_max = kwargs.get("multiturn_num_max", 5)
+        self.max_source_length = kwargs.get("max_source_length", 3000)
+        self.max_target_length = kwargs.get("max_target_length", 1024)
+        self.do_think = kwargs.get("do_think", True)
+        self.sys_prompt = kwargs.get("sys_prompt", True)
+
+        # used for dynamic output alignment
+        self.use_dynamic_output_ratio = kwargs.get("use_dynamic_output_ratio", 0.0)
+        self.min_output_mask_token_len = kwargs.get("min_mask_token_len", 1)
+        self.min_output_non_mask_token_len = kwargs.get("min_non_mask_token_len", 6)  # [eos]
+
+    def get_source_len(self, index):
+        """Get source len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_source_len(item)
+
+    def get_target_len(self, index):
+        """Get target len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_target_len(item)
+
+    def get_random_user_prompt(self, item, user_prompt):
+        """Get random user prompt.
+        
+            Args:
+                item: TODO.
+                user_prompt: TODO.
+            """
+        tasks = ["语音转写：", "Speech transcription:"]
+        language = item.get("language", None)
+        # LID in distill data is fake
+        language = None
+        if language is not None:
+            if language.lower() == "zh":
+                tasks.append("语音转写成中文：")
+                tasks.append("Transcribe speech into Chinese:")
+            elif language.lower() == "en":
+                tasks.append("语音转写成英文：")
+                tasks.append("Transcribe speech into English:")
+        if len(tasks) == 2:
+            task = random.choice(tasks)
+        elif len(tasks) == 4:
+            task = random.choices(tasks, weights=[0.4, 0.4, 0.1, 0.1])[0]
+        if "语音转写：<|startofspeech|>" in user_prompt:
+            user_prompt = user_prompt.replace("语音转写：<|startofspeech|>", task + "<|startofspeech|>")
+        elif "Speech transcription:<|startofspeech|>" in user_prompt:
+            user_prompt = user_prompt.replace("Speech transcription:<|startofspeech|>", task + "<|startofspeech|>")
+        return user_prompt
+
+    def __len__(self):
+        """Internal: len  ."""
+        return len(self.index_ds)
+
+    def __getitem__(self, index):
+        """Internal: getitem  .
+        
+            Args:
+                index: TODO.
+            """
+        output = None
+
+        for idx in range(self.retry):
+            if idx > 0:
+                logging.info(f"retry: {idx}")
+            badcase_flag = False
+            if idx == 0:
+                index_cur = index
+            else:
+                index_cur = torch.randint(0, len(self.index_ds), ()).item()
+
+            item = self.index_ds[index_cur]
+
+            system = item["system"]
+            user = item["user"]
+            assistant = item["assistant"]
+            is_noised = item.get("noised", False)
+            if len(user) < 1 or len(assistant) < 1:
+                logging.warning(f"item is error: {item}")
+                continue
+            input_ids, labels, fbank, fbank_lens, fbank_mask, fbank_beg, fake_token_len = (
+                [],
+                [],
+                [],
+                [],
+                [],
+                [],
+                [],
+            )
+
+            for i, (system_prompt, user_prompt, target_out) in enumerate(
+                    zip(system, user, assistant)
+            ):
+                if i >= self.multiturn_num_max:
+                    break
+                if len(input_ids) > self.max_token_length:
+                    logging.info(
+                        f"input_ids > max_token_length: {len(input_ids)}>{self.max_token_length}, {item}"
+                    )
+                    break
+
+                if self.prompt_classes is not None:
+                    asr_prompt = user_prompt.split("<|startofspeech|>")[0]
+                    language = self.prompt_classes.detect_language(asr_prompt)
+                    user_prompt_all_context = self.prompt_classes.get_prompt(item, language)
+                else:
+                    user_prompt_all_context = ""
+
+                if i == 0:
+                    source_input = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{user_prompt_all_context}{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
+                    if not self.sys_prompt:
+                        source_input = f"<|im_start|>user\n{user_prompt_all_context}{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
+                else:
+                    source_input = (
+                        f"<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
+                    )
+                if not self.do_think:
+                    source_input += "<think>\n\n</think>\n\n"
+                splits = self.pattern.split(source_input)
+                source_ids = []
+                fbank_i = []
+                fake_token_len_i = 0
+                fbank_beg_i = -1
+                fbank_lens_i = []
+                speech = []
+                speech_lengths = []
+                for k, sub_str in enumerate(splits):
+                    if not sub_str.startswith("<|startofspeech|>"):
+                        sub_token = self.tokenizer.encode(sub_str)
+                        source_ids += sub_token
+                    else:
+                        sub_str = sub_str.replace("<|startofspeech|>", "").replace(
+                            "<|endofspeech|>", ""
+                        )
+                        if sub_str.startswith("!"):
+                            try:
+                                data_src = load_audio_text_image_video(sub_str[1:], fs=self.fs)
+                                if self.preprocessor_noise is not None and not is_noised:
+                                    try:
+                                        data_src = self.preprocessor_noise(data_src.numpy())
+                                    except Exception as e:
+                                        logging.error(f"Generate noise audio failed: {e}")
+
+                                speech, speech_lengths = extract_fbank(
+                                    data_src,
+                                    data_type=self.data_type,
+                                    frontend=self.frontend,
+                                    is_final=True,
+                                )  # speech: [b, T, d]
+                                if speech_lengths > self.max_source_length:
+                                    logging.info(
+                                        f"speech_lengths > max_source_length: {speech_lengths}>{self.max_source_length}, {item}"
+                                    )
+                                    badcase_flag = True
+                            except Exception as e:
+                                logging.warning(
+                                    f"Loading wav failed! {str(e)}, {traceback.format_exc()}\n{item}"
+                                )
+                                badcase_flag = True
+                                continue
+                            if True:
+                                olens = 1 + (speech_lengths[0].item() - 3 + 2 * 1) // 2
+                                olens = 1 + (olens - 3 + 2 * 1) // 2
+                                fake_token_len_i = (olens - 1) // 2 + 1
+                            else:
+                                fake_token_len_i = speech_lengths[0].item()
+                            fake_token = [0] * fake_token_len_i
+                            fbank_beg_i = len(source_ids)
+                            source_ids += fake_token
+
+                if badcase_flag:
+                    continue
+                if fbank_beg_i > 0:
+                    fbank_beg += [fbank_beg_i + len(input_ids)]
+                    fake_token_len += [fake_token_len_i]
+                else:
+                    fbank_beg += [-1]
+                    fake_token_len += [0]
+
+                if target_out is not None and any(
+                        isinstance(item, dict) and "prev_content" in item for item in target_out
+                ):
+                    prev_value = next(
+                        (
+                            item["prev_content"]
+                            for item in target_out
+                            if isinstance(item, dict) and "prev_content" in item
+                        ),
+                        None,
+                    )
+                    source_ids += self.tokenizer.encode(prev_value)
+                    source_mask = [-100] * len(source_ids)
+                    target_out = f"{target_out[0]}<|im_end|>"
+                else:
+                    source_mask = [-100] * len(source_ids)
+                    target_out = f"{target_out}<|im_end|>"
+                target_ids = self.tokenizer.encode(target_out)
+
+                if len(target_ids) > self.max_target_length:
+                    logging.info(
+                        f"text_length: {len(target_ids)} > {self.max_target_length}, drop it: {item}"
+                    )
+                #  simulate prev-token fixed output
+                target_labels = target_ids.copy()
+                if np.random.rand() < self.use_dynamic_output_ratio:
+                    max_len = len(target_labels)
+                    min_output_mask_token_len = min(self.min_output_mask_token_len, max_len)
+                    min_output_non_mask_token_len = min(self.min_output_non_mask_token_len, max_len)
+                    if max_len - min_output_non_mask_token_len > min_output_mask_token_len:
+                        end_index = np.random.randint(min_output_mask_token_len,
+                                                      max_len - min_output_non_mask_token_len)
+                    else:
+                        end_index = max_len - min_output_non_mask_token_len
+                    if end_index > 0:
+                        target_labels[:end_index] = [-100] * end_index
+
+                input_ids += source_ids + target_ids
+                labels += source_mask + target_labels
+                if len(speech) > 0:
+                    fbank.append(speech[0, :, :])
+                    fbank_lens.append(speech_lengths)
+            if badcase_flag:
+                continue
+
+            input_ids = torch.tensor(input_ids, dtype=torch.int64)  # [: self.max_token_length]
+            attention_mask = torch.tensor([1] * len(input_ids), dtype=torch.int32)
+            labels = torch.tensor(labels, dtype=torch.int64)  # [: self.max_token_length]
+
+            fbank_beg = torch.tensor(fbank_beg, dtype=torch.int32)
+            fake_token_len = torch.tensor(fake_token_len, dtype=torch.int32)
+
+            output = {
+                "fbank_beg": fbank_beg,
+                "fake_token_len": fake_token_len,
+                "input_ids": input_ids,
+                "attention_mask": attention_mask,
+                "labels_ids": labels,
+            }
+            output["item"] = item
+            if len(fbank) > 0:
+                output["speech"] = fbank
+                output["speech_lengths"] = fbank_lens
+            if len(input_ids) > self.max_token_length:
+                logging.warning(
+                    f"len(input_ids): {len(input_ids)} > max_token_length: {self.max_token_length}, item: {item}"
+                )
+                continue
+
+            break
+
+        return output
+
+    def collator(self, samples: list = None):
+
+        """Collator.
+        
+            Args:
+                samples: TODO.
+            """
+        for idx in range(self.retry):
+            badcase_flag = False
+
+            outputs = {}
+            for sample in samples:
+                if sample is None:
+                    continue
+                for key in sample.keys():
+                    if key not in outputs:
+                        outputs[key] = []
+                    if isinstance(sample[key], (list, tuple)):
+                        outputs[key].extend(sample[key])
+                    else:
+                        outputs[key].append(sample[key])
+
+            for key, data_list in outputs.items():
+                if isinstance(data_list[0], torch.Tensor):
+                    if data_list[0].dtype == torch.int64 or data_list[0].dtype == torch.int32:
+
+                        pad_value = self.int_pad_value
+                    else:
+                        pad_value = self.float_pad_value
+
+                    outputs[key] = torch.nn.utils.rnn.pad_sequence(
+                        data_list, batch_first=True, padding_value=pad_value
+                    )
+
+            if self.batch_type != "example":
+                b, t = outputs["input_ids"].shape
+                if b > 1 and b * t > self.batch_size_token_max:
+                    logging.info(
+                        f"Warning, {idx}th, b*t: {b}*{t}={b * t} > batch_size_sample_max: {self.batch_size_token_max}, drop last data"
+                    )
+                    samples = samples[:-1]
+                    continue
+
+            break
+
+        return outputs
+
+
+@tables.register("index_ds_classes", "FunASR")
+class FunASR(torch.utils.data.Dataset):  # torch.utils.data.Dataset
+
+    def __init__(self, path: str, **kwargs):
+        """Initialize FunASR.
+        
+            Args:
+                path: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+
+        self.max_source_length = kwargs.get("max_source_length", 8000)
+        self.min_source_length = kwargs.get("min_source_length", 10)
+        self.max_target_length = kwargs.get("max_target_length", 2048)
+        self.min_target_length = kwargs.get("min_target_length", 0)
+        # self.max_token_length = kwargs.get("max_token_length", 2200)+
+        audio_downsample_rate = int(kwargs.get("audio_downsample_rate", 8))
+
+        is_training = kwargs.get("is_training", True)
+        if not (path.endswith(".jsonl") or path.endswith(".json")):
+            # jsonl list file
+            data_split_num = kwargs.get("data_split_num", 1)
+            data_split_i = kwargs.get("data_split_i", 0)
+
+            if not is_training:
+                data_split_num = 1
+                data_split_i = 0
+            with open(path, encoding="utf-8") as fin:
+                file_list_all = fin.readlines()
+
+                num_per_slice = (len(file_list_all) - 1) // data_split_num + 1  # 16
+                file_list = file_list_all[
+                            data_split_i * num_per_slice: (data_split_i + 1) * num_per_slice
+                            ]
+                logging.info(
+                    f"is_training: {is_training}, data_split_num: {data_split_num}, data_split_i: {data_split_i}, \nfile_list: {file_list}, \nfile_list_all: {file_list_all}"
+                )
+
+        else:
+            file_list = [path]
+
+        contents = []
+        total_whrs = 0.0
+        total_token_for_llm_B = 0.0
+        for file_json in file_list:
+            with open(file_json.strip(), encoding="utf-8") as fin:
+                for line in fin:
+                    try:
+                        data_dict = json.loads(line.strip())
+                    except Exception as e:
+                        logging.error(
+                            f"drop it, json error: {e}, line: {line}, file_json: {file_json}"
+                        )
+                        continue
+
+                    data = data_dict["messages"]
+                    if isinstance(data_dict.get("speech_length", 0), (list, tuple)):
+                        speech_length = int(data_dict.get("speech_length", [0])[0])
+                        text_length = int(data_dict.get("text_length", 0)[0])
+                    else:
+                        speech_length = int(data_dict.get("speech_length", 0))
+                        text_length = int(data_dict.get("text_length", 0))
+                    speech_length = int(speech_length)
+                    text_length = int(text_length)
+                    if speech_length > 0 and speech_length < 1:
+                        continue
+                    if text_length < 1:
+                        logging.warning(
+                            f"speech_length: {speech_length}, text_length: {text_length}, data: {data}, file_json: {file_json}"
+                        )
+                        if len(data) > 2:
+                            text_length = len(data[2]['content'])
+                        continue
+                    if speech_length > self.max_source_length:
+                        continue
+                    if speech_length < self.min_source_length:
+                        continue
+                    if text_length > self.max_target_length:
+                        continue
+
+                    system, user, assistant = [], [], []
+                    for i, item in enumerate(data):
+                        try:
+                            role = item["role"]
+                            content = item["content"]
+                        except KeyError:
+                            logging.error(
+                                f"drop it, KeyError: {item}, file_json: {file_json}"
+                            )
+                            continue
+
+                        if role == "system":
+                            system.append(content)
+                        elif role == "user":
+                            user.append(content)
+                        elif role == "assistant":
+                            if "prev_content" in item:
+                                prev_content = item["prev_content"]
+                                assistant.append([content, {"prev_content": prev_content}])
+                            else:
+                                assistant.append(content)
+                    if len(system) == 0:
+                        system = ["You are a helpful assistant."]
+                    system = system * len(user)
+
+                    contents_i = {
+                        "system": system,
+                        "user": user,
+                        "assistant": assistant,
+                        "source_len": speech_length + text_length,
+                    }
+                    if "key" in data_dict:
+                        contents_i["key"] = data_dict["key"] if not isinstance(data_dict.get("key", "key_01234"),
+                                                                               (list, tuple)) else data_dict["key"][0]
+
+                    if "hist_context" in data_dict:
+                        contents_i["hist_context"] = data_dict["hist_context"]
+                    if "hotwords" in data_dict:
+                        contents_i["hotwords"] = data_dict["hotwords"]
+                    if "asr_hotwords" in data_dict:
+                        contents_i["asr_hotwords"] = data_dict["asr_hotwords"]
+                    if "vad_segs" in data_dict:
+                        contents_i["vad_segs"] = data_dict["vad_segs"]
+                    if "word_list" in data_dict:
+                        contents_i["word_list"] = data_dict["word_list"]
+                    if "one_pass_result" in data_dict:
+                        contents_i["one_pass_result"] = data_dict["one_pass_result"]
+                    if "one_pass_wer" in data_dict:
+                        contents_i["one_pass_wer"] = data_dict["one_pass_wer"]
+                    if "noised" in data_dict:
+                        contents_i["noised"] = data_dict["noised"]
+
+                    if kwargs.get("save_meta", False):
+                        contents_i["meta"] = data_dict
+
+                    total_whrs += speech_length / 100.0 / 3600 / 10000 * audio_downsample_rate
+                    total_token_for_llm_B += (text_length + speech_length / 8) / 1000 / 1000 / 1000
+                    contents.append(contents_i)
+
+        self.contents = contents
+
+        logging.info(
+            f"\n\ntotal_num of samplers: {len(self.contents)}, total_whrs: {total_whrs:.5f}, total_token_for_llm_B: {total_token_for_llm_B:.5g}, {path}, {file_list}\n\n")
+
+    def __len__(self):
+        """Internal: len  ."""
+        return len(self.contents)
+
+    def __getitem__(self, index):
+
+        """Internal: getitem  .
+        
+            Args:
+                index: TODO.
+            """
+        data = self.contents[index]
+
+        return data
+
+    def get_source_len(self, data_dict):
+        """Get source len.
+        
+            Args:
+                data_dict: TODO.
+            """
+        source_len = data_dict.get("source_len", -1)
+        if source_len < 0:
+            source_len = len(data_dict["system"]) + len(data_dict["user"])
+        return source_len
+
+    def get_target_len(self, data_dict):
+
+        """Get target len.
+        
+            Args:
+                data_dict: TODO.
+            """
+        return 0

funasr/datasets/fun_asr_datasets/multicontext_prompt.py

@@ -0,0 +1,371 @@
+import numpy as np
+from funasr.register import tables
+import logging
+import random
+import re
+
+
+@tables.register("prompt_classes", "MultiContextPrompt")
+class MultiContextPrompt:
+    CONTEXT_TEMPLATES = {
+        'en': {
+            'header': "Please combine the context information provided below to complete the speech transcription task more accurately. If there is no relevant information, we will leave it blank.\n",
+            'fields': {
+                'hist_context': "Historical transcription: {hist_context}\n",
+                'one_pass_result': "One-pass result: {one_pass_result}\n",
+                'hotwords': "Hotword list: {hotwords}\n"
+            }
+        },
+        'zh': {
+            'header': "请结合下面提供的上下文信息，更加准确地完成语音转写任务。如果没有相关信息，我们会留空。\n",
+            'fields': {
+                'hist_context': "历史转写结果：{hist_context}\n",
+                'one_pass_result': "一遍解码结果：{one_pass_result}\n",
+                'hotwords': "热词列表：{hotwords}\n"
+            }
+        }
+    }
+
+    def __init__(self,
+                 use_hist=True,
+                 use_one_pass_result=True,
+                 use_hotwords=True,
+                 use_asr_hotwords=True,
+                 use_multi_lingual_prompt=True,
+                 **kwargs):
+        """Initialize MultiContextPrompt.
+        
+            Args:
+                use_hist: TODO.
+                use_one_pass_result: TODO.
+                use_hotwords: TODO.
+                use_asr_hotwords: TODO.
+                use_multi_lingual_prompt: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        self.use_hist = use_hist
+        self.use_one_pass_result = use_one_pass_result
+        self.use_hotwords = use_hotwords
+        self.use_asr_hotwords = use_asr_hotwords
+        self.use_multi_lingual_prompt = use_multi_lingual_prompt
+        self.kwargs = kwargs
+
+        chinese_hotwords_list = kwargs.get("chinese_hotwords_list", "")
+        english_hotwords_list = kwargs.get("english_hotwords_list", "")
+        if chinese_hotwords_list:
+            self.chinese_hotwords_list, self.chinese_hotwords_num = self.get_hotwords_list(chinese_hotwords_list)
+        else:
+            self.chinese_hotwords_list = None
+            self.chinese_hotwords_num = 0
+        logging.info(f"chinese_hotwords_num: {self.chinese_hotwords_num}")
+
+        if english_hotwords_list:
+            self.english_hotwords_list, self.english_hotwords_num = self.get_hotwords_list(english_hotwords_list)
+        else:
+            self.english_hotwords_list = None
+            self.english_hotwords_num = 0
+        logging.info(f"english_hotwords_num: {self.english_hotwords_num}")
+
+        self.max_neg_hotwords_num = kwargs.get("max_neg_hotwords_num", 900)
+        self.min_neg_hotwords_num = kwargs.get("min_neg_hotwords_num", 0)
+
+    def get_hotwords_list(self, hotwords_file):
+        """Get hotwords list.
+        
+            Args:
+                hotwords_file: TODO.
+            """
+        with open(hotwords_file, "r") as f:
+            hotwords_list = f.read().strip().split("\n")
+        return hotwords_list, len(hotwords_list)
+
+    def detect_language(self, text):
+        """Detect language.
+        
+            Args:
+                text: Text tensor or string input.
+            """
+        if isinstance(text, list):
+            text = " ".join(text)
+
+        chinese_pattern = re.compile(
+            "["
+            "\u4e00-\u9fff"  # CJK Unified Ideographs
+            "]+"
+        )
+
+        english_pattern = re.compile(r'[A-Za-z]+')
+
+        chinese_matches = chinese_pattern.findall(text)
+        english_matches = english_pattern.findall(text)
+
+        chinese_length = sum(len(match) for match in chinese_matches)
+        english_length = sum(len(match) for match in english_matches)
+
+        total_length = len(text)
+
+        if total_length == 0:
+            return 'zh'
+
+        if (chinese_length > english_length) and (chinese_length / total_length > 0.3):
+            return 'zh'
+        else:
+            return 'en'
+
+    def hotwords_sampling(self, hotwords):
+
+        # hotwords_list = hotwords.split(", ")
+        """Hotwords sampling.
+        
+            Args:
+                hotwords: TODO.
+            """
+        hotwords_list = hotwords
+        selected_hotwords = []
+        if self.max_neg_hotwords_num > -1:
+            max_neg_hotwords_num = min(self.max_neg_hotwords_num, len(hotwords_list))
+        else:
+            max_neg_hotwords_num = len(hotwords_list)
+
+        if self.min_neg_hotwords_num < max_neg_hotwords_num:
+            selected_hotwords_num = np.random.randint(self.min_neg_hotwords_num, max_neg_hotwords_num + 1)
+        else:
+            selected_hotwords_num = max_neg_hotwords_num
+        if selected_hotwords_num > 0:
+            selected_hotwords = np.random.choice(hotwords_list, selected_hotwords_num, replace=False).tolist()
+
+        return selected_hotwords, selected_hotwords_num
+
+    def get_prompt(self, item, language):
+        """Get prompt.
+        
+            Args:
+                item: TODO.
+                language: Language identifier.
+            """
+        template = self.CONTEXT_TEMPLATES[language]
+
+        prompt = template['header']
+
+        context_lines = []
+
+        if self.use_hist and item.get("hist_context"):
+            context_lines.append(template['fields']['hist_context'].format(hist_context=item["hist_context"]))
+
+        if self.use_one_pass_result and item.get("one_pass_result"):
+            context_lines.append(template['fields']['one_pass_result'].format(one_pass_result=item["one_pass_result"]))
+
+        hotwords = None
+        if self.use_hotwords and item.get("hotwords"):
+            hotwords = item["hotwords"]
+        if self.use_asr_hotwords and item.get("asr_hotwords"):
+            hotwords = item["asr_hotwords"]
+        if hotwords is not None and hotwords != "":
+            language = self.detect_language(hotwords)
+            if language == 'en':
+                neg_hotwords = self.english_hotwords_list
+            else:
+                neg_hotwords = self.chinese_hotwords_list
+            if neg_hotwords is not None:
+                selected_neg_hotwords, selected_neg_hotwords_num = self.hotwords_sampling(neg_hotwords)
+            else:
+                selected_neg_hotwords = []
+
+            if not isinstance(hotwords, list):
+                pos_hotwords = hotwords.split(", ")
+            else:
+                pos_hotwords = hotwords
+            hotwords = pos_hotwords + selected_neg_hotwords
+            random.shuffle(hotwords)
+            hotwords = ", ".join(hotwords)
+            context_lines.append(template['fields']['hotwords'].format(hotwords=hotwords))
+
+        if context_lines:
+            prompt += ''.join(context_lines)
+        else:
+            prompt += "\n\n\n"
+
+        return prompt
+
+    def get_inference_prompt(self, item, language="zh"):
+        """Get inference prompt.
+        
+            Args:
+                item: TODO.
+                language: Language identifier.
+            """
+        template = self.CONTEXT_TEMPLATES[language]
+
+        prompt = template['header']
+
+        context_lines = []
+
+        if self.use_hist and item.get("hist_context"):
+            context_lines.append(template['fields']['hist_context'].format(hist_context=item["hist_context"]))
+
+        if self.use_one_pass_result and item.get("one_pass_result"):
+            context_lines.append(template['fields']['one_pass_result'].format(one_pass_result=item["one_pass_result"]))
+
+        hotwords = None
+        if self.use_hotwords and item.get("hotwords"):
+            hotwords = item["hotwords"]
+        if self.use_asr_hotwords and item.get("asr_hotwords"):
+            hotwords = item["asr_hotwords"]
+        if hotwords is not None and hotwords != "":
+            print(f"hotwords: {hotwords}")
+            language = self.detect_language(hotwords)
+            if language == 'en':
+                neg_hotwords = self.english_hotwords_list
+            else:
+                neg_hotwords = self.chinese_hotwords_list
+            if neg_hotwords is not None:
+                selected_neg_hotwords, selected_neg_hotwords_num = self.hotwords_sampling(neg_hotwords)
+            else:
+                selected_neg_hotwords = []
+
+            if not isinstance(hotwords, list):
+                pos_hotwords = hotwords.split(", ")
+            else:
+                pos_hotwords = hotwords
+            hotwords = pos_hotwords + selected_neg_hotwords
+            print(f"selected_neg_hotwords_num: {selected_neg_hotwords_num}")
+            random.shuffle(hotwords)
+            hotwords = ", ".join(hotwords)
+            context_lines.append(template['fields']['hotwords'].format(hotwords=hotwords))
+
+        if context_lines:
+            prompt += ''.join(context_lines)
+        else:
+            prompt += "\n\n\n"
+
+        return prompt
+
+
+@tables.register("prompt_classes", "MultiContextPromptNew")
+class MultiContextPromptNew:
+    CONTEXT_TEMPLATES = {
+        'en': {
+            'header': "Please combine the context information to complete the speech transcription task more accurately. If there is no relevant information, we will leave it blank.\n\n",
+            'context_header': "**Context:**\n",
+            'fields': {
+                'hist_context': "Historical transcription: {hist_context}\n",
+                'one_pass_result': "One-pass result: {one_pass_result}\n",
+                'hotwords': "Hotword list: {hotwords}\n"
+            }
+        },
+        'zh': {
+            'header': "请结合上下文信息，更加准确地完成语音转写任务。如果没有相关信息，我们会留空。\n\n",
+            'context_header': "**上下文信息：**\n",
+            'fields': {
+                'hist_context': "历史转写结果：{hist_context}\n",
+                'one_pass_result': "一遍解码结果：{one_pass_result}\n",
+                'hotwords': "热词列表：{hotwords}\n"
+            }
+        }
+    }
+
+    def __init__(self,
+                 use_hist=True,
+                 use_one_pass_result=True,
+                 use_hotwords=True,
+                 use_multi_lingual_prompt=True,
+                 **kwargs):
+        """Initialize MultiContextPromptNew.
+        
+            Args:
+                use_hist: TODO.
+                use_one_pass_result: TODO.
+                use_hotwords: TODO.
+                use_multi_lingual_prompt: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        self.use_hist = use_hist
+        self.use_one_pass_result = use_one_pass_result
+        self.use_hotwords = use_hotwords
+        self.use_multi_lingual_prompt = use_multi_lingual_prompt
+
+        self.use_full_hotwords_ratio = kwargs.get("use_full_hotwords_ratio", 0.2)
+        self.max_hotwords_num = kwargs.get("max_hotwords_num", -1)
+        self.min_hotwords_num = kwargs.get("min_hotwords_num", 15)
+
+    def hotwords_sampling(self, hotwords):
+
+        """Hotwords sampling.
+        
+            Args:
+                hotwords: TODO.
+            """
+        hotwords_list = hotwords.split(", ")
+        if self.max_hotwords_num > 0:
+            max_hotwords_num = min(self.max_hotwords_num, len(hotwords_list))
+        else:
+            max_hotwords_num = len(hotwords_list)
+
+        if self.min_hotwords_num < max_hotwords_num:
+            selected_hotwords_num = np.random.randint(self.min_hotwords_num, max_hotwords_num + 1)
+        else:
+            selected_hotwords_num = max_hotwords_num
+
+        selected_hotwords = np.random.choice(hotwords_list, selected_hotwords_num, replace=False)
+        hotwords_list = ", ".join(selected_hotwords)
+
+        return hotwords_list, selected_hotwords_num
+
+    def get_prompt(self, item, language):
+        """Get prompt.
+        
+            Args:
+                item: TODO.
+                language: Language identifier.
+            """
+        template = self.CONTEXT_TEMPLATES[language]
+
+        prompt = template['header']
+
+        context_lines = []
+
+        if self.use_hist and item.get("hist_context"):
+            context_lines.append(template['fields']['hist_context'].format(hist_context=item["hist_context"]))
+
+        if self.use_one_pass_result and item.get("one_pass_result"):
+            context_lines.append(template['fields']['one_pass_result'].format(one_pass_result=item["one_pass_result"]))
+
+        if self.use_hotwords and item.get("hotwords"):
+            hotwords = item["hotwords"]
+            if np.random.rand() < self.use_full_hotwords_ratio:
+                hotwords = hotwords
+            else:
+                hotwords, selected_hotwords_num = self.hotwords_sampling(hotwords)
+            context_lines.append(template['fields']['hotwords'].format(hotwords=hotwords))
+
+        if context_lines:
+            prompt += template['context_header'] + ''.join(context_lines)
+
+        return prompt
+
+    def get_inference_prompt(self, hist_context="", one_pass_result="", hotwords=""):
+        """Get inference prompt.
+        
+            Args:
+                hist_context: TODO.
+                one_pass_result: TODO.
+                hotwords: TODO.
+            """
+        language = 'zh' if self.use_multi_lingual_prompt and np.random.rand() < 0.5 else 'en'
+        template = self.CONTEXT_TEMPLATES[language]
+
+        prompt = template['header']
+
+        context_lines = []
+
+        if hist_context:
+            context_lines.append(template['fields']['hist_context'].format(hist_context=hist_context))
+        if one_pass_result:
+            context_lines.append(template['fields']['one_pass_result'].format(one_pass_result=one_pass_result))
+        if hotwords:
+            context_lines.append(template['fields']['hotwords'].format(hotwords=hotwords))
+
+        if context_lines:
+            prompt += template['context_header'] + ''.join(context_lines)
+
+        return prompt

funasr/datasets/kws_datasets/datasets.py

@@ -0,0 +1,165 @@
+import torch
+import random
+
+
+from funasr.register import tables
+from funasr.utils.load_utils import extract_fbank, load_audio_text_image_video
+
+
+@tables.register("dataset_classes", "KwsMTDataset")
+class KwsMTDataset(torch.utils.data.Dataset):
+    """
+    KwsMTDataset, support multi tokenizers
+    """
+    def __init__(self,
+                 path,
+                 index_ds: str = None,
+                 frontend=None,
+                 tokenizer=None,
+                 is_training: bool = True,
+                 int_pad_value: int = -1,
+                 float_pad_value: float = 0.0,
+                 **kwargs,
+    ):
+        """Initialize KwsMTDataset.
+        
+            Args:
+                path: TODO.
+                index_ds: TODO.
+                frontend: Audio frontend for feature extraction.
+                tokenizer: Tokenizer instance for text encoding/decoding.
+                is_training: Boolean flag for training.
+                int_pad_value: TODO.
+                float_pad_value: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        index_ds_class = tables.index_ds_classes.get(index_ds)
+        self.index_ds = index_ds_class(path, **kwargs)
+
+        self.preprocessor_speech = None
+        self.preprocessor_text = None
+
+        if is_training:
+            preprocessor_speech = kwargs.get("preprocessor_speech", None)
+            if preprocessor_speech:
+                preprocessor_speech_class = tables.preprocessor_classes.get(preprocessor_speech)
+                preprocessor_speech = preprocessor_speech_class(
+                    **kwargs.get("preprocessor_speech_conf")
+                )
+            self.preprocessor_speech = preprocessor_speech
+
+            preprocessor_text = kwargs.get("preprocessor_text", None)
+            if preprocessor_text:
+                preprocessor_text_class = tables.preprocessor_classes.get(preprocessor_text)
+                preprocessor_text = preprocessor_text_class(**kwargs.get("preprocessor_text_conf"))
+            self.preprocessor_text = preprocessor_text
+
+        self.frontend = frontend
+        self.fs = 16000 if frontend is None else frontend.fs
+        self.data_type = "sound"
+        print(tokenizer)
+        self.tokenizer = tokenizer
+
+        self.int_pad_value = int_pad_value
+        self.float_pad_value = float_pad_value
+
+    def get_source_len(self, index):
+        """Get source len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_source_len(item)
+
+    def get_target_len(self, index):
+        """Get target len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_target_len(item)
+
+    def __len__(self):
+        """Internal: len  ."""
+        return len(self.index_ds)
+
+    def __getitem__(self, index):
+        """Internal: getitem  .
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        # import pdb;
+        # pdb.set_trace()
+        source = item["source"]
+        data_src = load_audio_text_image_video(source, fs=self.fs)
+        if self.preprocessor_speech:
+            data_src = self.preprocessor_speech(data_src, fs=self.fs)
+        speech, speech_lengths = extract_fbank(
+            data_src, data_type=self.data_type, frontend=self.frontend, is_final=True
+        ) # speech: [b, T, d]
+
+        target = item["target"]
+        if self.preprocessor_text:
+            target = self.preprocessor_text(target)
+
+        if self.tokenizer[0]:
+            ids = self.tokenizer[0].encode(target)
+            text = torch.tensor(ids, dtype=torch.int64)
+            # print("target: ", target,  ", ids: ", str(ids))
+        else:
+            ids = target
+            text = ids
+
+        if self.tokenizer[1]:
+            ids2 = self.tokenizer[1].encode(target)
+            text2 = torch.tensor(ids2, dtype=torch.int64)
+            # print("target: ", target,  ", ids2: ", str(ids2))
+        else:
+            ids2 = target
+            text2 = ids2
+
+        ids_lengths = len(ids)
+        text_lengths = torch.tensor([ids_lengths], dtype=torch.int32)
+
+        ids2_lengths = len(ids2)
+        text2_lengths = torch.tensor([ids2_lengths], dtype=torch.int32)
+
+        return {"speech": speech[0, :, :],
+                "speech_lengths": speech_lengths,
+                "text": text,
+                "text_lengths": text_lengths,
+                "text2": text2,
+                "text2_lengths": text2_lengths,
+                }
+
+
+    def collator(self, samples: list=None):
+        """Collator.
+        
+            Args:
+                samples: TODO.
+            """
+        outputs = {}
+        for sample in samples:
+            for key in sample.keys():
+                if key not in outputs:
+                    outputs[key] = []
+                outputs[key].append(sample[key])
+
+        for key, data_list in outputs.items():
+            if isinstance(data_list[0], torch.Tensor):
+                if data_list[0].dtype == torch.int64 or data_list[0].dtype == torch.int32:
+
+                    pad_value = self.int_pad_value
+                else:
+                    pad_value = self.float_pad_value
+
+                outputs[key] = torch.nn.utils.rnn.pad_sequence(
+                    data_list, batch_first=True, padding_value=pad_value
+                )
+        return outputs

funasr/datasets/llm_datasets/datasets.py

@@ -0,0 +1,530 @@
+import torch
+import copy
+
+from funasr.register import tables
+from funasr.utils.load_utils import extract_fbank, load_audio_text_image_video
+
+
+@tables.register("dataset_classes", "AudioLLMNARDataset")
+class AudioLLMNARDataset(torch.utils.data.Dataset):
+    """
+    AudioLLMDataset
+    """
+
+    def __init__(
+        self,
+        path,
+        index_ds: str = None,
+        frontend=None,
+        tokenizer=None,
+        int_pad_value: int = -1,
+        float_pad_value: float = 0.0,
+        **kwargs
+    ):
+        """Initialize AudioLLMNARDataset.
+        
+            Args:
+                path: TODO.
+                index_ds: TODO.
+                frontend: Audio frontend for feature extraction.
+                tokenizer: Tokenizer instance for text encoding/decoding.
+                int_pad_value: TODO.
+                float_pad_value: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        index_ds_class = tables.index_ds_classes.get(index_ds)
+        self.index_ds = index_ds_class(path, **kwargs)
+        preprocessor_speech = kwargs.get("preprocessor_speech", None)
+        if preprocessor_speech:
+            preprocessor_speech_class = tables.preprocessor_classes.get(preprocessor_speech)
+            preprocessor_speech = preprocessor_speech_class(
+                **kwargs.get("preprocessor_speech_conf", {})
+            )
+        self.preprocessor_speech = preprocessor_speech
+        preprocessor_text = kwargs.get("preprocessor_text", None)
+        if preprocessor_text:
+            preprocessor_text_class = tables.preprocessor_classes.get(preprocessor_text)
+            preprocessor_text = preprocessor_text_class(**kwargs.get("preprocessor_text_conf", {}))
+        self.preprocessor_text = preprocessor_text
+
+        self.frontend = frontend
+        self.fs = 16000 if frontend is None else frontend.fs
+        self.data_type = "sound"
+        self.tokenizer = tokenizer
+
+        self.float_pad_value = float_pad_value
+        self.prompt = kwargs.get("prompt", "Please copy the following text.")
+        self.prompt_pre = "USER: \nINSTRUCTION: {}\nINPUT: ".format(
+            self.prompt
+        )  # "USER: \nINSTRUCTION: {}\nINPUT: {}\nASSISTANT: "
+        self.prompt_af = ""
+        self.IGNORE_INDEX = kwargs.get("IGNORE_INDEX", -100)
+        self.int_pad_value = self.IGNORE_INDEX
+
+    def get_source_len(self, index):
+        """Get source len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_source_len(item)
+
+    def get_target_len(self, index):
+        """Get target len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_target_len(item)
+
+    def __len__(self):
+        """Internal: len  ."""
+        return len(self.index_ds)
+
+    def __getitem__(self, index):
+        """Internal: getitem  .
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        source = item["source"]
+        data_src = load_audio_text_image_video(source, fs=self.fs)
+        if self.preprocessor_speech:
+            data_src = self.preprocessor_speech(data_src, fs=self.fs)
+        speech, speech_lengths = extract_fbank(
+            data_src, data_type=self.data_type, frontend=self.frontend, is_final=True
+        )  # speech: [b, T, d]
+        speech = speech.squeeze(0)
+
+        target = item["target"]
+        if self.preprocessor_text:
+            target = self.preprocessor_text(target)
+
+        prompt_ids_pre = self.tokenizer.encode(self.prompt_pre)  # [bos,prompt]
+        prompt_ids_length = len(prompt_ids_pre)
+
+        # bos prompt audio bos target
+        # prompt_input = "{}{}".format(self.prompt_pre, target)
+        # prompt_input_ids = self.tokenizer.encode(prompt_input) #[bos, prompt, input]
+        # audio_length = len(prompt_input_ids) - prompt_ids_length
+        target_ids = self.tokenizer.encode(target)
+        if target_ids[0] == self.tokenizer.bos_token_id:
+            target_ids = target_ids[1:]
+        target_ids_length = len(target_ids)
+        audio_length = target_ids_length
+        input_ids = (
+            prompt_ids_pre + target_ids + [self.tokenizer.pad_token_id] + target_ids
+        )  # [bos, prompt, input, pad, target]
+        input_ids = torch.tensor(
+            copy.deepcopy(input_ids), dtype=torch.int64
+        )  # [bos, prompt, input, pad, target]
+        input_ids[prompt_ids_length : prompt_ids_length + audio_length] = (
+            -1
+        )  # [bos, prompt,-1, pad, target] # it is no need, only for check
+        attention_mask = input_ids.ge(-1)  # [true, true, true, true, true], length mask
+
+        # bos prompt audio target eos
+        # prompt_answer = "{}{}".format(self.prompt_pre, target)
+        # prompt_answer_ids = self.tokenizer.encode(prompt_answer) #[bos, prompt, input]
+        # answer_length = len(prompt_answer_ids) - prompt_ids_length
+        target_ids = self.tokenizer.encode(target)
+        if target_ids[0] == self.tokenizer.bos_token_id:
+            target_ids = target_ids[1:]
+        # target_ids_length = len(target_ids)
+        labels_ids = (
+            prompt_ids_pre + target_ids + target_ids + [self.tokenizer.eos_token_id]
+        )  # [bos, prompt, input, target, eos]
+        labels_ids = torch.tensor(
+            copy.deepcopy(labels_ids), dtype=torch.int64
+        )  # [bos, prompt, input, target, eos]
+        labels_ids[:prompt_ids_length] = -1  # [-1, -1, input, target, eos]
+        label_mask = labels_ids.ge(0)  # [false, false, true, true, true], length mask
+        labels_ids[~label_mask] = self.IGNORE_INDEX  # [-1, -1, input, target, eos]
+
+        audio_mask = (
+            [0] * prompt_ids_length + [1] * audio_length + [0] * target_ids_length + [0]
+        )  # [0, 0, 1, 0, 0]
+        audio_mask = torch.tensor(audio_mask, dtype=torch.float32)
+
+        ids = target_ids  # self.tokenizer.encode(target) # token ids is different from labels_ids
+        text = torch.tensor(ids, dtype=torch.int64)
+        text_lengths = torch.tensor([len(ids)], dtype=torch.int32)
+
+        prompt_bos_length = torch.tensor([len(prompt_ids_pre)], dtype=torch.int32)
+
+        return {
+            "speech": speech,
+            "speech_lengths": speech_lengths,
+            "text": text,
+            "text_lengths": text_lengths,
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels_ids": labels_ids,
+            "label_mask": label_mask,
+            "audio_mask": audio_mask,
+            "prompt_bos_length": prompt_bos_length,
+        }
+
+    def collator(self, samples: list = None):
+        """Collator.
+        
+            Args:
+                samples: TODO.
+            """
+        outputs = {}
+        for sample in samples:
+            for key in sample.keys():
+                if key not in outputs:
+                    outputs[key] = []
+                outputs[key].append(sample[key])
+
+        for key, data_list in outputs.items():
+            if isinstance(data_list[0], torch.Tensor):
+                if data_list[0].dtype == torch.int64 or data_list[0].dtype == torch.int32:
+
+                    pad_value = self.int_pad_value
+                else:
+                    pad_value = self.float_pad_value
+
+                outputs[key] = torch.nn.utils.rnn.pad_sequence(
+                    data_list, batch_first=True, padding_value=pad_value
+                )
+        return outputs
+
+
+@tables.register("dataset_classes", "AudioLLMDataset")
+class AudioLLMDataset(torch.utils.data.Dataset):
+    """
+    AudioLLMDataset
+    """
+
+    def __init__(
+        self,
+        path,
+        index_ds: str = None,
+        frontend=None,
+        tokenizer=None,
+        int_pad_value: int = -1,
+        float_pad_value: float = 0.0,
+        **kwargs
+    ):
+        """Initialize AudioLLMDataset.
+        
+            Args:
+                path: TODO.
+                index_ds: TODO.
+                frontend: Audio frontend for feature extraction.
+                tokenizer: Tokenizer instance for text encoding/decoding.
+                int_pad_value: TODO.
+                float_pad_value: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        index_ds_class = tables.index_ds_classes.get(index_ds)
+        self.index_ds = index_ds_class(path, **kwargs)
+        preprocessor_speech = kwargs.get("preprocessor_speech", None)
+        if preprocessor_speech:
+            preprocessor_speech_class = tables.preprocessor_classes.get(preprocessor_speech)
+            preprocessor_speech = preprocessor_speech_class(
+                **kwargs.get("preprocessor_speech_conf", {})
+            )
+        self.preprocessor_speech = preprocessor_speech
+        preprocessor_text = kwargs.get("preprocessor_text", None)
+        if preprocessor_text:
+            preprocessor_text_class = tables.preprocessor_classes.get(preprocessor_text)
+            preprocessor_text = preprocessor_text_class(**kwargs.get("preprocessor_text_conf", {}))
+        self.preprocessor_text = preprocessor_text
+
+        self.frontend = frontend
+        self.fs = 16000 if frontend is None else frontend.fs
+        self.data_type = "sound"
+        self.tokenizer = tokenizer
+
+        self.float_pad_value = float_pad_value
+        self.prompt = kwargs.get("prompt", "Transcribe speech to text.")
+        self.prompt_pre = "USER: \nINSTRUCTION: {}\nINPUT: ".format(
+            self.prompt
+        )  # "USER: \nINSTRUCTION: {}\nnINPUT: {}\nASSISTANT: "
+        self.prompt_af = ""
+        self.IGNORE_INDEX = kwargs.get("IGNORE_INDEX", -100)
+        self.int_pad_value = self.IGNORE_INDEX
+
+    def get_source_len(self, index):
+        """Get source len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_source_len(item)
+
+    def get_target_len(self, index):
+        """Get target len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_target_len(item)
+
+    def __len__(self):
+        """Internal: len  ."""
+        return len(self.index_ds)
+
+    def __getitem__(self, index):
+        """Internal: getitem  .
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        # import pdb;
+        # pdb.set_trace()
+        source = item["source"]
+        data_src = load_audio_text_image_video(source, fs=self.fs)
+        if self.preprocessor_speech:
+            data_src = self.preprocessor_speech(data_src, fs=self.fs)
+        speech, speech_lengths = extract_fbank(
+            data_src, data_type=self.data_type, frontend=self.frontend, is_final=True
+        )  # speech: [b, T, d]
+        speech = speech.squeeze(0)
+
+        target = item["target"]
+        if self.preprocessor_text:
+            target = self.preprocessor_text(target)
+
+        prompt_ids_pre = self.tokenizer.encode(self.prompt_pre)  # [bos,prompt]
+        prompt_ids_length = len(prompt_ids_pre)
+
+        prompt_input = "{}{}".format(self.prompt_pre, target)
+        prompt_input_ids = self.tokenizer.encode(prompt_input)
+        audio_length = len(prompt_input_ids) - prompt_ids_length
+        input_ids = prompt_input_ids + [self.tokenizer.pad_token_id]
+        input_ids = torch.tensor(input_ids, dtype=torch.int64)  # [bos, prompt, input, pad]
+        input_ids[prompt_ids_length:] = -1  # [bos, prompt,-1,-1]
+        attention_mask = input_ids.ge(-1)  # [true, true, true, true], length mask
+
+        prompt_answer = "{}{}".format(self.prompt_pre, target)
+        prompt_answer_ids = self.tokenizer.encode(prompt_answer)
+        answer_length = len(prompt_answer_ids) - prompt_ids_length
+        labels_ids = copy.deepcopy(prompt_input_ids) + [self.tokenizer.eos_token_id]
+        labels_ids = torch.tensor(labels_ids, dtype=torch.int64)  # [bos, prompt, input, eos]
+        labels_ids[:prompt_ids_length] = -1  # [-1, -1, input, eos]
+        label_mask = labels_ids.ge(0)  # [False,False,True,True]
+        labels_ids[~label_mask] = self.IGNORE_INDEX  # [-100,-100,input,eos]
+
+        audio_mask = [0] * prompt_ids_length + [1] * audio_length + [0]
+        audio_mask = torch.tensor(audio_mask, dtype=torch.float32)
+
+        ids = self.tokenizer.encode(target)  # token ids is different from labels_ids
+        text = torch.tensor(ids, dtype=torch.int64)
+        text_lengths = torch.tensor([len(ids)], dtype=torch.int32)
+
+        return {
+            "speech": speech,
+            "speech_lengths": speech_lengths,
+            "text": text,
+            "text_lengths": text_lengths,
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels_ids": labels_ids,
+            "label_mask": label_mask,
+            "audio_mask": audio_mask,
+        }
+
+    def collator(self, samples: list = None):
+        """Collator.
+        
+            Args:
+                samples: TODO.
+            """
+        outputs = {}
+        for sample in samples:
+            for key in sample.keys():
+                if key not in outputs:
+                    outputs[key] = []
+                outputs[key].append(sample[key])
+
+        for key, data_list in outputs.items():
+            if isinstance(data_list[0], torch.Tensor):
+                if data_list[0].dtype == torch.int64 or data_list[0].dtype == torch.int32:
+
+                    pad_value = self.int_pad_value
+                else:
+                    pad_value = self.float_pad_value
+
+                outputs[key] = torch.nn.utils.rnn.pad_sequence(
+                    data_list, batch_first=True, padding_value=pad_value
+                )
+        return outputs
+
+
+@tables.register("dataset_classes", "AudioLLMARDataset")
+class AudioLLMARDataset(torch.utils.data.Dataset):
+    """
+    AudioLLMDataset
+    """
+
+    def __init__(
+        self,
+        path,
+        index_ds: str = None,
+        frontend=None,
+        tokenizer=None,
+        int_pad_value: int = -1,
+        float_pad_value: float = 0.0,
+        **kwargs
+    ):
+        """Initialize AudioLLMARDataset.
+        
+            Args:
+                path: TODO.
+                index_ds: TODO.
+                frontend: Audio frontend for feature extraction.
+                tokenizer: Tokenizer instance for text encoding/decoding.
+                int_pad_value: TODO.
+                float_pad_value: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        index_ds_class = tables.index_ds_classes.get(index_ds)
+        self.index_ds = index_ds_class(path, **kwargs)
+        preprocessor_speech = kwargs.get("preprocessor_speech", None)
+        if preprocessor_speech:
+            preprocessor_speech_class = tables.preprocessor_classes.get(preprocessor_speech)
+            preprocessor_speech = preprocessor_speech_class(
+                **kwargs.get("preprocessor_speech_conf", {})
+            )
+        self.preprocessor_speech = preprocessor_speech
+        preprocessor_text = kwargs.get("preprocessor_text", None)
+        if preprocessor_text:
+            preprocessor_text_class = tables.preprocessor_classes.get(preprocessor_text)
+            preprocessor_text = preprocessor_text_class(**kwargs.get("preprocessor_text_conf", {}))
+        self.preprocessor_text = preprocessor_text
+
+        self.frontend = frontend
+        self.fs = 16000 if frontend is None else frontend.fs
+        self.data_type = "sound"
+        self.tokenizer = tokenizer
+
+        self.float_pad_value = float_pad_value
+        self.prompt = kwargs.get("prompt", "Transcribe speech to text.")
+        self.prompt_pre = "USER: \nINSTRUCTION: {}\nINPUT: ".format(
+            self.prompt
+        )  # "USER: \nINSTRUCTION: {}\nnINPUT: {}\nASSISTANT: "
+        self.prompt_af = ""
+        self.IGNORE_INDEX = kwargs.get("IGNORE_INDEX", -100)
+        self.int_pad_value = self.IGNORE_INDEX
+
+    def get_source_len(self, index):
+        """Get source len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_source_len(item)
+
+    def get_target_len(self, index):
+        """Get target len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_target_len(item)
+
+    def __len__(self):
+        """Internal: len  ."""
+        return len(self.index_ds)
+
+    def __getitem__(self, index):
+        """Internal: getitem  .
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        # import pdb;
+        # pdb.set_trace()
+        source = item["source"]
+        data_src = load_audio_text_image_video(source, fs=self.fs)
+        if self.preprocessor_speech:
+            data_src = self.preprocessor_speech(data_src, fs=self.fs)
+        speech, speech_lengths = extract_fbank(
+            data_src, data_type=self.data_type, frontend=self.frontend, is_final=True
+        )  # speech: [b, T, d]
+        speech = speech.squeeze(0)
+
+        target = item["target"]
+        if self.preprocessor_text:
+            target = self.preprocessor_text(target)
+
+        prompt_ids_pre = self.tokenizer.encode(self.prompt_pre)  # [bos,prompt]
+        prompt_ids_length = len(prompt_ids_pre)
+
+        prompt_input = "{}{}".format(self.prompt_pre, target)
+        prompt_input_ids = self.tokenizer.encode(prompt_input)
+        audio_length = len(prompt_input_ids) - prompt_ids_length
+        input_ids = prompt_input_ids + [self.tokenizer.pad_token_id]
+        input_ids = torch.tensor(input_ids, dtype=torch.int64)  # [bos, prompt, input, pad]
+        input_ids[prompt_ids_length:] = -1  # [bos, prompt,-1,-1]
+        attention_mask = input_ids.ge(-1)  # [true, true, true, true], length mask
+
+        prompt_answer = "{}{}".format(self.prompt_pre, target)
+        prompt_answer_ids = self.tokenizer.encode(prompt_answer)
+        answer_length = len(prompt_answer_ids) - prompt_ids_length
+        labels_ids = copy.deepcopy(prompt_input_ids) + [self.tokenizer.eos_token_id]
+        labels_ids = torch.tensor(labels_ids, dtype=torch.int64)  # [bos, prompt, input, eos]
+        labels_ids[:prompt_ids_length] = -1  # [-1, -1, input, eos]
+        label_mask = labels_ids.ge(0)  # [False,False,True,True]
+        labels_ids[~label_mask] = self.IGNORE_INDEX  # [-100,-100,input,eos]
+
+        audio_mask = [0] * prompt_ids_length + [1] * audio_length + [0]
+        audio_mask = torch.tensor(audio_mask, dtype=torch.float32)
+
+        ids = self.tokenizer.encode(target)  # token ids is different from labels_ids
+        text = torch.tensor(ids, dtype=torch.int64)
+        text_lengths = torch.tensor([len(ids)], dtype=torch.int32)
+
+        return {
+            "speech": speech,
+            "speech_lengths": speech_lengths,
+            "text": text,
+            "text_lengths": text_lengths,
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels_ids": labels_ids,
+            "label_mask": label_mask,
+            "audio_mask": audio_mask,
+        }
+
+    def collator(self, samples: list = None):
+        """Collator.
+        
+            Args:
+                samples: TODO.
+            """
+        outputs = {}
+        for sample in samples:
+            for key in sample.keys():
+                if key not in outputs:
+                    outputs[key] = []
+                outputs[key].append(sample[key])
+
+        for key, data_list in outputs.items():
+            if isinstance(data_list[0], torch.Tensor):
+                if data_list[0].dtype == torch.int64 or data_list[0].dtype == torch.int32:
+
+                    pad_value = self.int_pad_value
+                else:
+                    pad_value = self.float_pad_value
+
+                outputs[key] = torch.nn.utils.rnn.pad_sequence(
+                    data_list, batch_first=True, padding_value=pad_value
+                )
+        return outputs

funasr/datasets/llm_datasets/preprocessor.py

@@ -0,0 +1,45 @@
+import os
+import json
+import torch
+import logging
+import concurrent.futures
+import librosa
+import torch.distributed as dist
+from typing import Collection
+import torch
+import torchaudio
+from torch import nn
+import random
+import re
+import string
+from funasr.tokenizer.cleaner import TextCleaner
+from funasr.register import tables
+
+
+@tables.register("preprocessor_classes", "TextPreprocessRemovePunctuation")
+class TextPreprocessRemovePunctuation(nn.Module):
+    def __init__(self, **kwargs):
+        """Initialize TextPreprocessRemovePunctuation.
+        
+            Args:
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+
+    def forward(self, text, **kwargs):
+        # 定义英文标点符号
+        """Forward pass for training.
+        
+            Args:
+                text: Text tensor or string input.
+                **kwargs: Additional keyword arguments.
+            """
+        en_punct = string.punctuation
+        # 定义中文标点符号（部分常用的）
+        cn_punct = "。？！，、；：“”‘’（）《》【】…—～·"
+        # 合并英文和中文标点符号
+        all_punct = en_punct + cn_punct
+        # 创建正则表达式模式，匹配任何在all_punct中的字符
+        punct_pattern = re.compile("[{}]".format(re.escape(all_punct)))
+        # 使用正则表达式的sub方法替换掉这些字符
+        return punct_pattern.sub("", text)

funasr/datasets/llm_datasets_qwenaudio/datasets.py

@@ -0,0 +1,191 @@
+import torch
+import copy
+
+from funasr.register import tables
+from funasr.utils.load_utils import extract_fbank, load_audio_text_image_video
+
+
+@tables.register("dataset_classes", "AudioLLMQwenAudioDataset")
+class AudioLLMQwenAudioDataset(torch.utils.data.Dataset):
+    """
+    AudioLLMDataset
+    """
+
+    def __init__(
+        self,
+        path,
+        index_ds: str = None,
+        frontend=None,
+        tokenizer=None,
+        int_pad_value: int = -1,
+        float_pad_value: float = 0.0,
+        **kwargs
+    ):
+        """Initialize AudioLLMQwenAudioDataset.
+        
+            Args:
+                path: TODO.
+                index_ds: TODO.
+                frontend: Audio frontend for feature extraction.
+                tokenizer: Tokenizer instance for text encoding/decoding.
+                int_pad_value: TODO.
+                float_pad_value: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        index_ds_class = tables.index_ds_classes.get(index_ds)
+        self.index_ds = index_ds_class(path, **kwargs)
+        preprocessor_speech = kwargs.get("preprocessor_speech", None)
+        if preprocessor_speech:
+            preprocessor_speech_class = tables.preprocessor_classes.get(preprocessor_speech)
+            preprocessor_speech = preprocessor_speech_class(
+                **kwargs.get("preprocessor_speech_conf", {})
+            )
+        self.preprocessor_speech = preprocessor_speech
+        preprocessor_text = kwargs.get("preprocessor_text", None)
+        if preprocessor_text:
+            preprocessor_text_class = tables.preprocessor_classes.get(preprocessor_text)
+            preprocessor_text = preprocessor_text_class(**kwargs.get("preprocessor_text_conf", {}))
+        self.preprocessor_text = preprocessor_text
+
+        self.frontend = frontend
+        self.fs = 16000 if frontend is None else frontend.fs
+        self.data_type = "sound"
+        self.tokenizer = tokenizer
+
+        self.float_pad_value = float_pad_value
+        self.prompt = kwargs.get("prompt", "Transcribe speech to text.")
+        # self.prompt_pre = "USER: \nINSTRUCTION: {}\nINPUT: ".format(self.prompt)  # "USER: \nINSTRUCTION: {}\nnINPUT: {}\nASSISTANT: "
+        self.prompt_af = ""
+        self.IGNORE_INDEX = kwargs.get("IGNORE_INDEX", -100)
+        self.int_pad_value = self.IGNORE_INDEX
+        self.audio_adaptor_downsample_rate = kwargs.get("audio_adaptor_downsample_rate", 5)
+        self.audio_encoder_downsample_rate = kwargs.get("audio_encoder_downsample_rate", 2)
+        self.prompt_template = "{}"
+        self.answer_template = "{}"
+
+    def get_source_len(self, index):
+        """Get source len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_source_len(item)
+
+    def get_target_len(self, index):
+        """Get target len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_target_len(item)
+
+    def __len__(self):
+        """Internal: len  ."""
+        return len(self.index_ds)
+
+    def __getitem__(self, index):
+        """Internal: getitem  .
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        source = item["source"]
+        data_src = load_audio_text_image_video(source, fs=self.fs)
+        if self.preprocessor_speech:
+            data_src = self.preprocessor_speech(data_src, fs=self.fs)
+        speech, speech_lengths = extract_fbank(
+            data_src, data_type=self.data_type, frontend=self.frontend, is_final=True
+        )  # speech: [b, T, d]
+        speech = speech.squeeze(0)
+
+        audio_pseudo_length = (
+            (speech.shape[0] + 1)
+            // self.audio_adaptor_downsample_rate
+            // self.audio_encoder_downsample_rate
+        )
+        audio_pseudo = torch.full((audio_pseudo_length,), -1)  # placeholder
+
+        target = item["target"]
+        if self.preprocessor_text:
+            target = self.preprocessor_text(target)
+
+        self.prompt_pre = self.prompt_template.format(self.prompt)
+        prompt_ids_pre = self.tokenizer.encode(self.prompt_pre)  # [bos,prompt]
+        prompt_pre_length = len(prompt_ids_pre)
+
+        # input
+        input = self.answer_template.format(target.lower())
+        prompt_input = "{}{}".format(self.prompt_pre, input)
+        prompt_input_ids = self.tokenizer.encode(prompt_input)  # [bos, prompt, input]
+        # audio_length = len(prompt_input_ids) - prompt_pre_length
+        input_ids = prompt_input_ids + [self.tokenizer.pad_token_id]  # [bos, prompt, input, pad]
+        input_ids_length = len(input_ids)
+        input_ids = torch.tensor(input_ids, dtype=torch.int64)  # [bos, prompt, input, pad]
+        input_ids = torch.cat((audio_pseudo, input_ids))  # [audio, bos, prompt, input, pad]
+        # input_ids[:audio_pseudo_length] = -1 # [-1, bos, prompt, input, pad]
+        attention_mask = input_ids.ge(-1)  # [true, true, true, true, true], length mask
+        # input_ids[prompt_pre_length:] = -1  # [bos, prompt,-1,-1]
+        # attention_mask = input_ids.ge(-1)  # [true, true, true, true], length mask
+
+        # label
+        answer = self.answer_template.format(target.lower())
+        prompt_answer = "{}{}".format(self.prompt_pre, answer)
+        prompt_answer_ids = self.tokenizer.encode(prompt_answer)
+        # answer_length = len(prompt_answer_ids) - prompt_pre_length
+        labels_ids = copy.deepcopy(prompt_answer_ids) + [self.tokenizer.eos_token_id]
+        labels_ids = torch.tensor(labels_ids, dtype=torch.int64)  # [bos, prompt, answer, eos]
+        labels_ids = torch.cat((audio_pseudo, labels_ids))  # [audio, bos, prompt, answer, eos]
+        labels_ids[: audio_pseudo_length + prompt_pre_length] = -1  # [-1, -1, -1, answer, eos]
+        # labels_ids[:prompt_pre_length] = -1  # [-1, -1, input, eos]
+        label_mask = labels_ids.ge(0)  # [false, false, false, true, true]
+        labels_ids[~label_mask] = self.IGNORE_INDEX  # [-100, -100, -100, answer, eos]
+
+        # audio_mask for input_ids
+        audio_mask = [1] * audio_pseudo_length + [0] * input_ids_length
+        audio_mask = torch.tensor(audio_mask, dtype=torch.float32)
+
+        ids = self.tokenizer.encode(target)  # token ids is different from labels_ids
+        text = torch.tensor(ids, dtype=torch.int64)
+        text_lengths = torch.tensor([len(ids)], dtype=torch.int32)
+
+        return {
+            "speech": speech,
+            "speech_lengths": speech_lengths,
+            "text": text,
+            "text_lengths": text_lengths,
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels_ids": labels_ids,
+            "label_mask": label_mask,
+            "audio_mask": audio_mask,
+        }
+
+    def collator(self, samples: list = None):
+        """Collator.
+        
+            Args:
+                samples: TODO.
+            """
+        outputs = {}
+        for sample in samples:
+            for key in sample.keys():
+                if key not in outputs:
+                    outputs[key] = []
+                outputs[key].append(sample[key])
+
+        for key, data_list in outputs.items():
+            if isinstance(data_list[0], torch.Tensor):
+                if data_list[0].dtype == torch.int64 or data_list[0].dtype == torch.int32:
+
+                    pad_value = self.int_pad_value
+                else:
+                    pad_value = self.float_pad_value
+
+                outputs[key] = torch.nn.utils.rnn.pad_sequence(
+                    data_list, batch_first=True, padding_value=pad_value
+                )
+        return outputs

funasr/datasets/llm_datasets_vicuna/datasets.py

@@ -0,0 +1,191 @@
+import torch
+import copy
+
+from funasr.register import tables
+from funasr.utils.load_utils import extract_fbank, load_audio_text_image_video
+
+
+@tables.register("dataset_classes", "AudioLLMVicunaDataset")
+class AudioLLMVicunaDataset(torch.utils.data.Dataset):
+    """
+    AudioLLMDataset
+    """
+
+    def __init__(
+        self,
+        path,
+        index_ds: str = None,
+        frontend=None,
+        tokenizer=None,
+        int_pad_value: int = -1,
+        float_pad_value: float = 0.0,
+        **kwargs
+    ):
+        """Initialize AudioLLMVicunaDataset.
+        
+            Args:
+                path: TODO.
+                index_ds: TODO.
+                frontend: Audio frontend for feature extraction.
+                tokenizer: Tokenizer instance for text encoding/decoding.
+                int_pad_value: TODO.
+                float_pad_value: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        index_ds_class = tables.index_ds_classes.get(index_ds)
+        self.index_ds = index_ds_class(path, **kwargs)
+        preprocessor_speech = kwargs.get("preprocessor_speech", None)
+        if preprocessor_speech:
+            preprocessor_speech_class = tables.preprocessor_classes.get(preprocessor_speech)
+            preprocessor_speech = preprocessor_speech_class(
+                **kwargs.get("preprocessor_speech_conf", {})
+            )
+        self.preprocessor_speech = preprocessor_speech
+        preprocessor_text = kwargs.get("preprocessor_text", None)
+        if preprocessor_text:
+            preprocessor_text_class = tables.preprocessor_classes.get(preprocessor_text)
+            preprocessor_text = preprocessor_text_class(**kwargs.get("preprocessor_text_conf", {}))
+        self.preprocessor_text = preprocessor_text
+
+        self.frontend = frontend
+        self.fs = 16000 if frontend is None else frontend.fs
+        self.data_type = "sound"
+        self.tokenizer = tokenizer
+
+        self.float_pad_value = float_pad_value
+        self.prompt = kwargs.get("prompt", "Transcribe speech to text.")
+        # self.prompt_pre = "USER: \nINSTRUCTION: {}\nINPUT: ".format(self.prompt)  # "USER: \nINSTRUCTION: {}\nnINPUT: {}\nASSISTANT: "
+        self.prompt_af = ""
+        self.IGNORE_INDEX = kwargs.get("IGNORE_INDEX", -100)
+        self.int_pad_value = self.IGNORE_INDEX
+        self.audio_adaptor_downsample_rate = kwargs.get("audio_adaptor_downsample_rate", 5)
+        self.audio_encoder_downsample_rate = kwargs.get("audio_encoder_downsample_rate", 2)
+        self.prompt_template = "USER: {}\n ASSISTANT:"
+        self.answer_template = "{}"
+
+    def get_source_len(self, index):
+        """Get source len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_source_len(item)
+
+    def get_target_len(self, index):
+        """Get target len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_target_len(item)
+
+    def __len__(self):
+        """Internal: len  ."""
+        return len(self.index_ds)
+
+    def __getitem__(self, index):
+        """Internal: getitem  .
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        source = item["source"]
+        data_src = load_audio_text_image_video(source, fs=self.fs)
+        if self.preprocessor_speech:
+            data_src = self.preprocessor_speech(data_src, fs=self.fs)
+        speech, speech_lengths = extract_fbank(
+            data_src, data_type=self.data_type, frontend=self.frontend, is_final=True
+        )  # speech: [b, T, d]
+        speech = speech.squeeze(0)
+
+        audio_pseudo_length = (
+            (speech.shape[0] + 1)
+            // self.audio_adaptor_downsample_rate
+            // self.audio_encoder_downsample_rate
+        )
+        audio_pseudo = torch.full((audio_pseudo_length,), -1)  # placeholder
+
+        target = item["target"]
+        if self.preprocessor_text:
+            target = self.preprocessor_text(target)
+
+        self.prompt_pre = self.prompt_template.format(self.prompt)
+        prompt_ids_pre = self.tokenizer.encode(self.prompt_pre)  # [bos,prompt]
+        prompt_pre_length = len(prompt_ids_pre)
+
+        # input
+        input = self.answer_template.format(target.lower())
+        prompt_input = "{}{}".format(self.prompt_pre, input)
+        prompt_input_ids = self.tokenizer.encode(prompt_input)  # [bos, prompt, input]
+        # audio_length = len(prompt_input_ids) - prompt_pre_length
+        input_ids = prompt_input_ids + [self.tokenizer.pad_token_id]  # [bos, prompt, input, pad]
+        input_ids_length = len(input_ids)
+        input_ids = torch.tensor(input_ids, dtype=torch.int64)  # [bos, prompt, input, pad]
+        input_ids = torch.cat((audio_pseudo, input_ids))  # [audio, bos, prompt, input, pad]
+        # input_ids[:audio_pseudo_length] = -1 # [-1, bos, prompt, input, pad]
+        attention_mask = input_ids.ge(-1)  # [true, true, true, true, true], length mask
+        # input_ids[prompt_pre_length:] = -1  # [bos, prompt,-1,-1]
+        # attention_mask = input_ids.ge(-1)  # [true, true, true, true], length mask
+
+        # label
+        answer = self.answer_template.format(target.lower())
+        prompt_answer = "{}{}".format(self.prompt_pre, answer)
+        prompt_answer_ids = self.tokenizer.encode(prompt_answer)
+        # answer_length = len(prompt_answer_ids) - prompt_pre_length
+        labels_ids = copy.deepcopy(prompt_answer_ids) + [self.tokenizer.eos_token_id]
+        labels_ids = torch.tensor(labels_ids, dtype=torch.int64)  # [bos, prompt, answer, eos]
+        labels_ids = torch.cat((audio_pseudo, labels_ids))  # [audio, bos, prompt, answer, eos]
+        labels_ids[: audio_pseudo_length + prompt_pre_length] = -1  # [-1, -1, -1, answer, eos]
+        # labels_ids[:prompt_pre_length] = -1  # [-1, -1, input, eos]
+        label_mask = labels_ids.ge(0)  # [false, false, false, true, true]
+        labels_ids[~label_mask] = self.IGNORE_INDEX  # [-100, -100, -100, answer, eos]
+
+        # audio_mask for input_ids
+        audio_mask = [1] * audio_pseudo_length + [0] * input_ids_length
+        audio_mask = torch.tensor(audio_mask, dtype=torch.float32)
+
+        ids = self.tokenizer.encode(target)  # token ids is different from labels_ids
+        text = torch.tensor(ids, dtype=torch.int64)
+        text_lengths = torch.tensor([len(ids)], dtype=torch.int32)
+
+        return {
+            "speech": speech,
+            "speech_lengths": speech_lengths,
+            "text": text,
+            "text_lengths": text_lengths,
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels_ids": labels_ids,
+            "label_mask": label_mask,
+            "audio_mask": audio_mask,
+        }
+
+    def collator(self, samples: list = None):
+        """Collator.
+        
+            Args:
+                samples: TODO.
+            """
+        outputs = {}
+        for sample in samples:
+            for key in sample.keys():
+                if key not in outputs:
+                    outputs[key] = []
+                outputs[key].append(sample[key])
+
+        for key, data_list in outputs.items():
+            if isinstance(data_list[0], torch.Tensor):
+                if data_list[0].dtype == torch.int64 or data_list[0].dtype == torch.int32:
+
+                    pad_value = self.int_pad_value
+                else:
+                    pad_value = self.float_pad_value
+
+                outputs[key] = torch.nn.utils.rnn.pad_sequence(
+                    data_list, batch_first=True, padding_value=pad_value
+                )
+        return outputs

funasr/datasets/openai_datasets/datasets.py

@@ -0,0 +1,545 @@
+import logging
+import re
+import torch
+import random
+import traceback
+from funasr.register import tables
+from funasr.utils.load_utils import extract_fbank, load_audio_text_image_video
+
+
+@tables.register("dataset_classes", "OpenAIDataset")
+class OpenAIDataset(torch.utils.data.Dataset):
+    """
+    SenseVoiceDataset
+    """
+
+    def __init__(
+        self,
+        path,
+        index_ds: str = None,
+        frontend=None,
+        tokenizer=None,
+        int_pad_value: int = -1,
+        float_pad_value: float = 0.0,
+        **kwargs,
+    ):
+        """Initialize OpenAIDataset.
+        
+            Args:
+                path: TODO.
+                index_ds: TODO.
+                frontend: Audio frontend for feature extraction.
+                tokenizer: Tokenizer instance for text encoding/decoding.
+                int_pad_value: TODO.
+                float_pad_value: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        index_ds_class = tables.index_ds_classes.get(index_ds)
+        self.index_ds = index_ds_class(path, **kwargs)
+        preprocessor_speech = kwargs.get("preprocessor_speech", None)
+        if preprocessor_speech:
+            preprocessor_speech_class = tables.preprocessor_classes.get(preprocessor_speech)
+            preprocessor_speech = preprocessor_speech_class(
+                **kwargs.get("preprocessor_speech_conf")
+            )
+        self.preprocessor_speech = preprocessor_speech
+        preprocessor_text = kwargs.get("preprocessor_text", None)
+        if preprocessor_text:
+            preprocessor_text_class = tables.preprocessor_classes.get(preprocessor_text)
+            preprocessor_text = preprocessor_text_class(**kwargs.get("preprocessor_text_conf"))
+        self.preprocessor_text = preprocessor_text
+
+        self.frontend = frontend
+        self.fs = 16000 if frontend is None else frontend.fs
+        self.data_type = "sound"
+        self.tokenizer = tokenizer
+
+        self.int_pad_value = int_pad_value
+        self.float_pad_value = float_pad_value
+        self.sos = kwargs.get("sos", "<|startoftranscript|>")
+        self.eos = kwargs.get("eos", "<|endoftext|>")
+        self.batch_size = kwargs.get("batch_size")
+        self.batch_type = kwargs.get("batch_type")
+        self.prompt_ids_len = 0
+        self.retry = kwargs.get("retry", 100)
+
+        self.permute = False
+        from funasr.frontends.whisper_frontend import WhisperFrontend
+
+        if isinstance(self.frontend, WhisperFrontend):
+            self.permute = True
+
+        self.pattern = re.compile(r"(<\|startofspeech\|>.*?<\|endofspeech\|>)")
+        # self.kwargs = kwargs
+        self.max_token_length = kwargs.get("max_token_length", 1024)
+        self.batch_size_scale_ratio_max = kwargs.get("batch_size_scale_ratio_max", 1.5)
+        self.batch_size_token_max = kwargs.get("batch_size_token_max", 2500)
+        self.audio_adaptor_downsample_rate = kwargs.get("audio_adaptor_downsample_rate", 2)
+        self.audio_encoder_downsample_rate = kwargs.get("audio_encoder_downsample_rate", 4)
+
+    def get_source_len(self, index):
+        """Get source len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_source_len(item)
+
+    def get_target_len(self, index):
+        """Get target len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_target_len(item)
+
+    def __len__(self):
+        """Internal: len  ."""
+        return len(self.index_ds)
+
+    def __getitem__(self, index):
+        # import pdb;
+        # pdb.set_trace()
+
+        """Internal: getitem  .
+        
+            Args:
+                index: TODO.
+            """
+        output = None
+
+        for idx in range(self.retry):
+            badcase_flag = False
+            if idx == 0:
+                index_cur = index
+            else:
+                index_cur = torch.randint(0, len(self.index_ds), ()).item()
+
+            item = self.index_ds[index_cur]
+
+            system = item["system"]
+            user = item["user"]
+            assistant = item["assistant"]
+
+            input_ids, labels, fbank, fbank_lens, fbank_mask, fbank_beg = [], [], [], [], [], []
+
+            for i, (system_prompt, user_prompt, target_out) in enumerate(
+                zip(system, user, assistant)
+            ):
+
+                source_input = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
+
+                splits = self.pattern.split(source_input)
+                source_ids = []
+                fbank_mask_i = []
+                fbank_beg_i = []
+                fbank_lens_i = []
+                for k, sub_str in enumerate(splits):
+                    if not sub_str.startswith("<|startofspeech|>"):
+                        sub_token = self.tokenizer.encode(sub_str)
+                        source_ids += sub_token
+                        fbank_mask_i += [0] * len(sub_token)
+                    else:
+                        sub_str = sub_str.replace("<|startofspeech|>", "").replace(
+                            "<|endofspeech|>", ""
+                        )
+                        if sub_str.startswith("!"):
+                            try:
+                                data_src = load_audio_text_image_video(sub_str[1:], fs=self.fs)
+                            except Exception as e:
+                                logging.error(
+                                    f"Loading wav failed! {str(e)}, {traceback.format_exc()}"
+                                )
+                                badcase_flag = True
+                                continue
+                            speech, speech_lengths = extract_fbank(
+                                data_src,
+                                data_type=self.data_type,
+                                frontend=self.frontend,
+                                is_final=True,
+                            )  # speech: [b, T, d]
+                            if self.permute:
+                                speech = speech.permute(0, 2, 1)
+                            # if speech_lengths > self.batch_size:
+                            #     continue
+                            if self.audio_encoder_downsample_rate == 4:
+                                olens = 1 + (speech_lengths[0].item() - 3 + 2 * 1) // 2
+                                olens = 1 + (olens - 3 + 2 * 1) // 2
+                            elif self.audio_encoder_downsample_rate == 1:
+                                olens = speech_lengths[0].item()
+
+                            sub_token_len = (olens - 1) // self.audio_adaptor_downsample_rate + 1
+                            sub_token = [0] * sub_token_len
+                            fbank_beg_i = [len(source_ids)]
+                            source_ids += sub_token
+                            fbank_mask_i += [1] * len(sub_token)
+
+                if badcase_flag:
+                    continue
+                source_mask = [-100] * len(source_ids)
+                target_out = f"{target_out}<|im_end|>"
+                target_ids = self.tokenizer.encode(target_out)
+                input_ids += source_ids + target_ids
+                labels += source_mask + target_ids
+                fbank_mask += fbank_mask_i
+                fbank_beg.append(fbank_beg_i)
+
+            if len(input_ids) > self.max_token_length:
+                logging.info(
+                    f"input_ids > max_token_length: {len(input_ids)}>{self.max_token_length}, {item}"
+                )
+                badcase_flag = True
+            if badcase_flag:
+                continue
+            input_ids = torch.tensor(input_ids, dtype=torch.int64)  # [: self.max_token_length]
+            attention_mask = torch.tensor([1] * len(input_ids), dtype=torch.int32)
+            labels = torch.tensor(labels, dtype=torch.int64)  # [: self.max_token_length]
+
+            fbank = speech[0, :, :]
+            fbank_lens = speech_lengths
+            fbank_mask = torch.tensor(fbank_mask, dtype=torch.float32)
+            fbank_beg = torch.tensor(fbank_beg, dtype=torch.int32)
+
+            output = {
+                "speech": fbank,
+                "speech_lengths": fbank_lens,
+                "fbank_mask": fbank_mask,
+                "fbank_beg": fbank_beg,
+                "input_ids": input_ids,
+                "attention_mask": attention_mask,
+                "labels_ids": labels,
+            }
+            break
+
+        return output
+
+    def collator(self, samples: list = None):
+
+        """Collator.
+        
+            Args:
+                samples: TODO.
+            """
+        for idx in range(self.retry):
+            badcase_flag = False
+
+            outputs = {}
+            for sample in samples:
+                if sample is None:
+                    continue
+                for key in sample.keys():
+                    if key not in outputs:
+                        outputs[key] = []
+                    outputs[key].append(sample[key])
+
+            for key, data_list in outputs.items():
+                if isinstance(data_list[0], torch.Tensor):
+                    if data_list[0].dtype == torch.int64 or data_list[0].dtype == torch.int32:
+
+                        pad_value = self.int_pad_value
+                    else:
+                        pad_value = self.float_pad_value
+
+                    outputs[key] = torch.nn.utils.rnn.pad_sequence(
+                        data_list, batch_first=True, padding_value=pad_value
+                    )
+
+            if self.batch_type != "example":
+                b, t = outputs["input_ids"].shape
+                if b > 1 and b * t > self.batch_size_token_max:
+                    logging.info(
+                        f"Warning, {idx}th, b*t: {b}*{t}={b * t} > batch_size_sample_max: {self.batch_size_token_max}, drop last data"
+                    )
+                    samples = samples[:-1]
+                    continue
+
+            break
+
+        return outputs
+
+
+@tables.register("dataset_classes", "OpenAIDatasetMultiTurn")
+class OpenAIDatasetMultiTurn(torch.utils.data.Dataset):
+    """
+    SenseVoiceDataset
+    """
+
+    def __init__(
+        self,
+        path,
+        index_ds: str = None,
+        frontend=None,
+        tokenizer=None,
+        int_pad_value: int = -1,
+        float_pad_value: float = 0.0,
+        **kwargs,
+    ):
+        """Initialize OpenAIDatasetMultiTurn.
+        
+            Args:
+                path: TODO.
+                index_ds: TODO.
+                frontend: Audio frontend for feature extraction.
+                tokenizer: Tokenizer instance for text encoding/decoding.
+                int_pad_value: TODO.
+                float_pad_value: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        index_ds_class = tables.index_ds_classes.get(index_ds)
+        self.index_ds = index_ds_class(path, **kwargs)
+        preprocessor_speech = kwargs.get("preprocessor_speech", None)
+        if preprocessor_speech:
+            preprocessor_speech_class = tables.preprocessor_classes.get(preprocessor_speech)
+            preprocessor_speech = preprocessor_speech_class(
+                **kwargs.get("preprocessor_speech_conf")
+            )
+        self.preprocessor_speech = preprocessor_speech
+        preprocessor_text = kwargs.get("preprocessor_text", None)
+        if preprocessor_text:
+            preprocessor_text_class = tables.preprocessor_classes.get(preprocessor_text)
+            preprocessor_text = preprocessor_text_class(**kwargs.get("preprocessor_text_conf"))
+        self.preprocessor_text = preprocessor_text
+
+        self.frontend = frontend
+        self.fs = 16000 if frontend is None else frontend.fs
+        self.data_type = "sound"
+        self.tokenizer = tokenizer
+
+        self.int_pad_value = int_pad_value
+        self.float_pad_value = float_pad_value
+        self.sos = kwargs.get("sos", "<|startoftranscript|>")
+        self.eos = kwargs.get("eos", "<|endoftext|>")
+        self.batch_size = kwargs.get("batch_size")
+        self.batch_type = kwargs.get("batch_type")
+        self.prompt_ids_len = 0
+        self.retry = kwargs.get("retry", 100)
+
+        self.permute = False
+        from funasr.frontends.whisper_frontend import WhisperFrontend
+
+        if isinstance(self.frontend, WhisperFrontend):
+            self.permute = True
+
+        self.pattern = re.compile(r"(<\|startofspeech\|>.*?<\|endofspeech\|>)")
+        # self.kwargs = kwargs
+        self.max_token_length = kwargs.get("max_token_length", 1500)
+        self.batch_size_scale_ratio_max = kwargs.get("batch_size_scale_ratio_max", 1.5)
+        self.batch_size_token_max = kwargs.get("batch_size_token_max", 2500)
+        self.multiturn_num_max = kwargs.get("multiturn_num_max", 5)
+        self.max_source_length = kwargs.get("max_source_length", 3000)
+
+    def get_source_len(self, index):
+        """Get source len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_source_len(item)
+
+    def get_target_len(self, index):
+        """Get target len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_target_len(item)
+
+    def __len__(self):
+        """Internal: len  ."""
+        return len(self.index_ds)
+
+    def __getitem__(self, index):
+        # import pdb
+        #
+        # pdb.set_trace()
+
+        """Internal: getitem  .
+        
+            Args:
+                index: TODO.
+            """
+        output = None
+
+        for idx in range(self.retry):
+            badcase_flag = False
+            if idx == 0:
+                index_cur = index
+            else:
+                index_cur = torch.randint(0, len(self.index_ds), ()).item()
+
+            item = self.index_ds[index_cur]
+
+            system = item["system"]
+            user = item["user"]
+            assistant = item["assistant"]
+
+            input_ids, labels, fbank, fbank_lens, fbank_mask, fbank_beg, fake_token_len = (
+                [],
+                [],
+                [],
+                [],
+                [],
+                [],
+                [],
+            )
+
+            for i, (system_prompt, user_prompt, target_out) in enumerate(
+                zip(system, user, assistant)
+            ):
+                if i >= self.multiturn_num_max:
+                    break
+                if len(input_ids) > self.max_token_length:
+                    logging.info(
+                        f"input_ids > max_token_length: {len(input_ids)}>{self.max_token_length}, {item}"
+                    )
+                    break
+
+                if i == 0:
+                    source_input = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
+                else:
+                    source_input = (
+                        f"<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
+                    )
+
+                splits = self.pattern.split(source_input)
+                source_ids = []
+                fbank_i = []
+                fbank_mask_i = []
+                fake_token_len_i = 0
+                fbank_beg_i = -1
+                fbank_lens_i = []
+                for k, sub_str in enumerate(splits):
+                    if not sub_str.startswith("<|startofspeech|>"):
+                        sub_token = self.tokenizer.encode(sub_str)
+                        source_ids += sub_token
+                        fbank_mask_i += [0] * len(sub_token)
+                    else:
+                        sub_str = sub_str.replace("<|startofspeech|>", "").replace(
+                            "<|endofspeech|>", ""
+                        )
+                        if sub_str.startswith("!"):
+                            try:
+                                data_src = load_audio_text_image_video(sub_str[1:], fs=self.fs)
+                            except Exception as e:
+                                logging.error(
+                                    f"Loading wav failed! {str(e)}, {traceback.format_exc()}"
+                                )
+                                badcase_flag = True
+                                continue
+                            speech, speech_lengths = extract_fbank(
+                                data_src,
+                                data_type=self.data_type,
+                                frontend=self.frontend,
+                                is_final=True,
+                            )  # speech: [b, T, d]
+                            if speech_lengths > self.max_source_length:
+                                logging.info(
+                                    f"speech_lengths > max_source_length: {speech_lengths}>{self.max_source_length}, {item}"
+                                )
+                                badcase_flag = True
+                            if self.permute:
+                                speech = speech.permute(0, 2, 1)
+                            # if speech_lengths > self.batch_size:
+                            #     continue
+
+                            olens = 1 + (speech_lengths[0].item() - 3 + 2 * 1) // 2
+                            olens = 1 + (olens - 3 + 2 * 1) // 2
+                            fake_token_len_i = (olens - 1) // 2 + 1
+                            fake_token = [0] * fake_token_len_i
+                            fbank_beg_i = len(source_ids)
+                            source_ids += fake_token
+                            fbank_mask_i += [1] * len(fake_token)
+
+                if badcase_flag:
+                    continue
+
+                fbank_beg += [fbank_beg_i + len(input_ids)]
+                fake_token_len += [fake_token_len_i]
+                source_mask = [-100] * len(source_ids)
+                target_out = f"{target_out}<|im_end|>"
+                target_ids = self.tokenizer.encode(target_out)
+                input_ids += source_ids + target_ids
+                labels += source_mask + target_ids
+                fbank.append(speech[0, :, :])
+                fbank_mask += fbank_mask_i
+                fbank_lens.append(speech_lengths)
+
+            if badcase_flag:
+                continue
+
+            input_ids = torch.tensor(input_ids, dtype=torch.int64)  # [: self.max_token_length]
+            attention_mask = torch.tensor([1] * len(input_ids), dtype=torch.int32)
+            labels = torch.tensor(labels, dtype=torch.int64)  # [: self.max_token_length]
+
+            # fbank = speech[0, :, :]
+            # fbank_lens = torch.tensor(fbank_lens, dtype=torch.int32)
+            fbank_mask = torch.tensor(fbank_mask, dtype=torch.float32)
+            fbank_beg = torch.tensor(fbank_beg, dtype=torch.int32)
+            fake_token_len = torch.tensor(fake_token_len, dtype=torch.int32)
+
+            output = {
+                "speech": fbank,
+                "speech_lengths": fbank_lens,
+                "fbank_mask": fbank_mask,
+                "fbank_beg": fbank_beg,
+                "fake_token_len": fake_token_len,
+                "input_ids": input_ids,
+                "attention_mask": attention_mask,
+                "labels_ids": labels,
+            }
+            break
+
+        return output
+
+    def collator(self, samples: list = None):
+
+        """Collator.
+        
+            Args:
+                samples: TODO.
+            """
+        for idx in range(self.retry):
+            badcase_flag = False
+
+            outputs = {}
+            for sample in samples:
+                if sample is None:
+                    continue
+                for key in sample.keys():
+                    if key not in outputs:
+                        outputs[key] = []
+                    if isinstance(sample[key], (list, tuple)):
+                        outputs[key].extend(sample[key])
+                    else:
+                        outputs[key].append(sample[key])
+
+            for key, data_list in outputs.items():
+                if isinstance(data_list[0], torch.Tensor):
+                    if data_list[0].dtype == torch.int64 or data_list[0].dtype == torch.int32:
+
+                        pad_value = self.int_pad_value
+                    else:
+                        pad_value = self.float_pad_value
+
+                    outputs[key] = torch.nn.utils.rnn.pad_sequence(
+                        data_list, batch_first=True, padding_value=pad_value
+                    )
+
+            if self.batch_type != "example":
+                b, t = outputs["input_ids"].shape
+                if b > 1 and b * t > self.batch_size_token_max:
+                    logging.info(
+                        f"Warning, {idx}th, b*t: {b}*{t}={b * t} > batch_size_sample_max: {self.batch_size_token_max}, drop last data"
+                    )
+                    samples = samples[:-1]
+                    continue
+
+            break
+
+        return outputs

funasr/datasets/openai_datasets/index_ds.py

@@ -0,0 +1,138 @@
+import os
+import json
+import torch
+import logging
+
+import librosa
+import random
+import torch.distributed as dist
+
+from funasr.register import tables
+
+
+@tables.register("index_ds_classes", "OpenAIIndexDSJsonl")
+class OpenAIIndexDSJsonl(torch.utils.data.Dataset):  # torch.utils.data.Dataset
+
+    def __init__(self, path: str, **kwargs):
+        """Initialize OpenAIIndexDSJsonl.
+        
+            Args:
+                path: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+
+        self.max_source_length = kwargs.get("max_source_length", 3000)
+        self.min_source_length = kwargs.get("min_source_length", 0)
+        self.max_target_length = kwargs.get("max_target_length", 2048)
+        self.min_target_length = kwargs.get("min_target_length", 0)
+        self.max_token_length = kwargs.get("max_token_length", 2200)
+
+        is_training = kwargs.get("is_training", True)
+        if not (path.endswith(".jsonl") or path.endswith(".json")):
+            # jsonl list file
+            data_split_num = kwargs.get("data_split_num", 1)
+            data_split_i = kwargs.get("data_split_i", 0)
+
+            if not is_training:
+                data_split_num = 1
+                data_split_i = 0
+            with open(path, encoding="utf-8") as fin:
+                file_list_all = fin.readlines()
+
+                num_per_slice = (len(file_list_all) - 1) // data_split_num + 1  # 16
+                file_list = file_list_all[
+                    data_split_i * num_per_slice : (data_split_i + 1) * num_per_slice
+                ]
+                logging.info(
+                    f"is_training: {is_training}, data_split_num: {data_split_num}, data_split_i: {data_split_i}, \nfile_list: {file_list}, \nfile_list_all: {file_list_all}"
+                )
+
+        else:
+            file_list = [path]
+
+        contents = []
+        for file_json in file_list:
+            with open(file_json.strip(), encoding="utf-8") as fin:
+                for line in fin:
+                    data_dict = json.loads(line.strip())
+                    data = data_dict["messages"]
+                    speech_length = data_dict.get("speech_length", -1) // 8
+                    text_length = data_dict.get("text_length", 0)
+                    if speech_length > self.max_source_length:
+                        logging.info(
+                            "speech_length: {speech_length} > {self.max_source_length}, drop it"
+                        )
+                        continue
+                    if text_length > self.max_target_length:
+                        continue
+
+                    self.max_target_length = kwargs.get("max_target_length", 2048)
+
+                    system, user, assistant = [], [], []
+                    for i, item in enumerate(data):
+                        role = item["role"]
+                        content = item["content"]
+                        if role == "system":
+                            system.append(content)
+                        elif role == "user":
+                            user.append(content)
+                        elif role == "assistant":
+                            assistant.append(content)
+
+                    system = system * len(user)
+
+                    contents_i = {
+                        "system": system,
+                        "user": user,
+                        "assistant": assistant,
+                        "source_len": speech_length + text_length,
+                    }
+                    contents.append(contents_i)
+
+        self.contents = contents
+
+        logging.info("total_num of samplers: {}, {}".format(len(self.contents), path))
+
+    def __len__(self):
+        """Internal: len  ."""
+        return len(self.contents)
+
+    def __getitem__(self, index):
+
+        """Internal: getitem  .
+        
+            Args:
+                index: TODO.
+            """
+        data = self.contents[index]
+
+        return data
+
+    def get_source_len(self, data_dict):
+        """Get source len.
+        
+            Args:
+                data_dict: TODO.
+            """
+        source_len = data_dict.get("source_len", -1)
+        if source_len < 0:
+            source_len = len(data_dict["system"]) + len(data_dict["user"])
+        return source_len
+
+    def get_target_len(self, data_dict):
+
+        """Get target len.
+        
+            Args:
+                data_dict: TODO.
+            """
+        return 0
+
+
+if __name__ == "__main__":
+    index_ds = OpenAIIndexDSJsonl(
+        path="/Users/zhifu/funasr1.0/test_local/data_tmp/tmp_wav_10.jsonl"
+    )
+    print(index_ds.contents)
+    pass

funasr/datasets/sense_voice_datasets/datasets.py

@@ -0,0 +1,506 @@
+import logging
+
+import re
+import torch
+import random
+import traceback
+from funasr.register import tables
+from funasr.utils.load_utils import extract_fbank, load_audio_text_image_video
+
+
+@tables.register("dataset_classes", "SenseVoiceDataset")
+class SenseVoiceDataset(torch.utils.data.Dataset):
+    """
+    SenseVoiceDataset
+    """
+
+    def __init__(
+        self,
+        path,
+        index_ds: str = None,
+        frontend=None,
+        tokenizer=None,
+        int_pad_value: int = -1,
+        float_pad_value: float = 0.0,
+        **kwargs,
+    ):
+        """Initialize SenseVoiceDataset.
+        
+            Args:
+                path: TODO.
+                index_ds: TODO.
+                frontend: Audio frontend for feature extraction.
+                tokenizer: Tokenizer instance for text encoding/decoding.
+                int_pad_value: TODO.
+                float_pad_value: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        index_ds_class = tables.index_ds_classes.get(index_ds)
+        self.index_ds = index_ds_class(path, **kwargs)
+        preprocessor_speech = kwargs.get("preprocessor_speech", None)
+        if preprocessor_speech:
+            preprocessor_speech_class = tables.preprocessor_classes.get(preprocessor_speech)
+            preprocessor_speech = preprocessor_speech_class(
+                **kwargs.get("preprocessor_speech_conf")
+            )
+        self.preprocessor_speech = preprocessor_speech
+        preprocessor_text = kwargs.get("preprocessor_text", None)
+        if preprocessor_text:
+            preprocessor_text_class = tables.preprocessor_classes.get(preprocessor_text)
+            preprocessor_text = preprocessor_text_class(**kwargs.get("preprocessor_text_conf"))
+        self.preprocessor_text = preprocessor_text
+
+        self.frontend = frontend
+        self.fs = 16000 if frontend is None else frontend.fs
+        self.data_type = "sound"
+        self.tokenizer = tokenizer
+
+        self.int_pad_value = int_pad_value
+        self.float_pad_value = float_pad_value
+        self.sos = kwargs.get("sos", "<|startoftranscript|>")
+        self.eos = kwargs.get("eos", "<|endoftext|>")
+        self.batch_size = kwargs.get("batch_size")
+        self.batch_type = kwargs.get("batch_type")
+        self.prompt_ids_len = 0
+        self.retry = kwargs.get("retry", 5)
+
+        self.permute = False
+        from funasr.frontends.whisper_frontend import WhisperFrontend
+
+        if isinstance(self.frontend, WhisperFrontend):
+            self.permute = True
+
+    def get_source_len(self, index):
+        """Get source len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_source_len(item)
+
+    def get_target_len(self, index):
+        """Get target len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_target_len(item)
+
+    def __len__(self):
+        """Internal: len  ."""
+        return len(self.index_ds)
+
+    def __getitem__(self, index):
+
+        """Internal: getitem  .
+        
+            Args:
+                index: TODO.
+            """
+        output = None
+        for idx in range(self.retry):
+            if idx == 0:
+                index_cur = index
+            else:
+                index_cur = torch.randint(0, len(self.index_ds), ()).item()
+
+            item = self.index_ds[index_cur]
+
+            source = item["source"]
+            try:
+                data_src = load_audio_text_image_video(source, fs=self.fs)
+            except Exception as e:
+                logging.error(f"Loading wav failed! {str(e)}, {traceback.format_exc()}")
+                continue
+
+            if self.preprocessor_speech:
+                data_src = self.preprocessor_speech(data_src, fs=self.fs)
+            speech, speech_lengths = extract_fbank(
+                data_src, data_type=self.data_type, frontend=self.frontend, is_final=True
+            )  # speech: [b, T, d]
+
+            if speech_lengths > self.batch_size:
+                continue
+            if self.permute:
+                speech = speech.permute(0, 2, 1)
+            target = item["target"]
+            if self.preprocessor_text:
+                target = self.preprocessor_text(target)
+
+            task = item.get("prompt", "<|ASR|>")
+            text_language = item.get("text_language", "<|zh|>")
+
+            if isinstance(self.sos, str):
+                prompt = f"{self.sos}{task}{text_language}"
+                prompt_ids = self.tokenizer.encode(prompt, allowed_special="all")
+            else:
+                prompt = f"{task}{text_language}"
+                prompt_ids = self.tokenizer.encode(prompt, allowed_special="all")
+                prompt_ids = [self.sos] + prompt_ids
+
+            prompt_ids_len = len(prompt_ids) - 1  # [sos, task]
+            self.prompt_ids_len = prompt_ids_len
+
+            target_ids = self.tokenizer.encode(target, allowed_special="all")
+            target_ids_len = len(target_ids) + 1  # [lid, text]
+            if target_ids_len > 200:
+                continue
+
+            if isinstance(self.eos, str):
+                eos = self.tokenizer.encode(self.eos, allowed_special="all")  # [eos]
+            else:
+                eos = [self.eos]
+
+            ids = prompt_ids + target_ids + eos  # [sos, task, lid, text, eos]
+            ids_lengths = len(ids)
+
+            text = torch.tensor(ids, dtype=torch.int64)
+            text_lengths = torch.tensor([ids_lengths], dtype=torch.int32)
+
+            target_mask = (
+                [0] * (prompt_ids_len) + [1] * (target_ids_len) + [1]
+            )  # [sos, task, lid, text, eos]: [0, 0, 1, 1, 1]
+            target_mask_lengths = len(target_mask)
+            target_mask = torch.tensor(target_mask, dtype=torch.float32)
+            target_mask_lengths = torch.tensor([target_mask_lengths], dtype=torch.int32)
+
+            output = {
+                "speech": speech[0, :, :],
+                "speech_lengths": speech_lengths,
+                "text": text,
+                "text_lengths": text_lengths,
+                "target_mask": target_mask,
+                "target_mask_lengths": target_mask_lengths,
+            }
+            break
+
+        return output
+
+    def collator(self, samples: list = None):
+        """Collator.
+        
+            Args:
+                samples: TODO.
+            """
+        outputs = {}
+        for sample in samples:
+            if sample is None:
+                continue
+            for key in sample.keys():
+                if key not in outputs:
+                    outputs[key] = []
+                outputs[key].append(sample[key])
+
+        if len(outputs) < 1:
+            logging.error(f"ERROR: data is empty!")
+            outputs = {
+                "speech": torch.rand((10, 128), dtype=torch.float32)[None, :, :],
+                "speech_lengths": torch.tensor(
+                    [
+                        10,
+                    ],
+                    dtype=torch.int32,
+                )[:, None],
+                "text": torch.tensor(
+                    [
+                        58836,
+                    ],
+                    dtype=torch.int32,
+                )[None, :],
+                "text_lengths": torch.tensor(
+                    [
+                        1,
+                    ],
+                    dtype=torch.int32,
+                )[:, None],
+                "target_mask": torch.tensor([[0] * (self.prompt_ids_len) + [1] * (1) + [1]])[
+                    None, :
+                ],
+            }
+            return outputs
+
+        for key, data_list in outputs.items():
+            if isinstance(data_list[0], torch.Tensor):
+                if data_list[0].dtype == torch.int64 or data_list[0].dtype == torch.int32:
+
+                    pad_value = self.int_pad_value
+                else:
+                    pad_value = self.float_pad_value
+
+                outputs[key] = torch.nn.utils.rnn.pad_sequence(
+                    data_list, batch_first=True, padding_value=pad_value
+                )
+
+        if self.batch_type != "example":
+            for i in range(10):
+                outputs = self._filter_badcase(outputs, i=i)
+
+        return outputs
+
+    def _filter_badcase(self, outputs, i=0):
+        """Internal: filter badcase.
+        
+            Args:
+                outputs: TODO.
+                i: TODO.
+            """
+        b, t, _ = outputs["speech"].shape
+
+        if b * t > self.batch_size * 1.25:
+            beg = torch.randint(0, 2, ()).item()
+            if b < 2:
+                beg = 0
+            logging.info(
+                f"Warning, b * t: {b * t} > {self.batch_size}, drop half data {i}th, beg:{beg}"
+            )
+            for key, data_list in outputs.items():
+                outputs[key] = outputs[key][beg : beg + b : 2]
+
+            speech_lengths_max = outputs["speech_lengths"].max().item()
+            outputs["speech"] = outputs["speech"][:, :speech_lengths_max, :]
+            text_lengths_max = outputs["text_lengths"].max().item()
+            outputs["text"] = outputs["text"][:, :text_lengths_max]
+            target_mask_lengths_max = outputs["target_mask_lengths"].max().item()
+            outputs["target_mask"] = outputs["target_mask"][:, :target_mask_lengths_max]
+
+        return outputs
+
+
+@tables.register("dataset_classes", "SenseVoiceCTCDataset")
+class SenseVoiceCTCDataset(torch.utils.data.Dataset):
+    """
+    SenseVoiceCTCDataset
+    """
+
+    def __init__(
+        self,
+        path,
+        index_ds: str = None,
+        frontend=None,
+        tokenizer=None,
+        int_pad_value: int = -1,
+        float_pad_value: float = 0.0,
+        **kwargs,
+    ):
+        """Initialize SenseVoiceCTCDataset.
+        
+            Args:
+                path: TODO.
+                index_ds: TODO.
+                frontend: Audio frontend for feature extraction.
+                tokenizer: Tokenizer instance for text encoding/decoding.
+                int_pad_value: TODO.
+                float_pad_value: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        index_ds_class = tables.index_ds_classes.get(index_ds)
+        self.index_ds = index_ds_class(path, **kwargs)
+        preprocessor_speech = kwargs.get("preprocessor_speech", None)
+        if preprocessor_speech:
+            preprocessor_speech_class = tables.preprocessor_classes.get(preprocessor_speech)
+            preprocessor_speech = preprocessor_speech_class(
+                **kwargs.get("preprocessor_speech_conf")
+            )
+        self.preprocessor_speech = preprocessor_speech
+        preprocessor_text = kwargs.get("preprocessor_text", None)
+        if preprocessor_text:
+            preprocessor_text_class = tables.preprocessor_classes.get(preprocessor_text)
+            preprocessor_text = preprocessor_text_class(**kwargs.get("preprocessor_text_conf"))
+        self.preprocessor_text = preprocessor_text
+
+        self.frontend = frontend
+        self.fs = 16000 if frontend is None else frontend.fs
+        self.data_type = "sound"
+        self.tokenizer = tokenizer
+
+        self.int_pad_value = int_pad_value
+        self.float_pad_value = float_pad_value
+        self.sos = kwargs.get("sos", "<|startoftranscript|>")
+        self.eos = kwargs.get("eos", "<|endoftext|>")
+        self.batch_size = kwargs.get("batch_size")
+        self.batch_type = kwargs.get("batch_type")
+        self.prompt_ids_len = 0
+        self.retry = kwargs.get("retry", 5)
+
+        self.permute = False
+        from funasr.frontends.whisper_frontend import WhisperFrontend
+
+        if isinstance(self.frontend, WhisperFrontend):
+            self.permute = True
+
+    def get_source_len(self, index):
+        """Get source len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_source_len(item)
+
+    def get_target_len(self, index):
+        """Get target len.
+        
+            Args:
+                index: TODO.
+            """
+        item = self.index_ds[index]
+        return self.index_ds.get_target_len(item)
+
+    def __len__(self):
+        """Internal: len  ."""
+        return len(self.index_ds)
+
+    def __getitem__(self, index):
+
+        """Internal: getitem  .
+        
+            Args:
+                index: TODO.
+            """
+        output = None
+        for idx in range(self.retry):
+            if idx == 0:
+                index_cur = index
+            else:
+                index_cur = torch.randint(0, len(self.index_ds), ()).item()
+
+            item = self.index_ds[index_cur]
+
+            source = item["source"]
+            try:
+                data_src = load_audio_text_image_video(source, fs=self.fs)
+            except Exception as e:
+                logging.error(f"Loading wav failed! {str(e)}, {traceback.format_exc()}")
+                continue
+
+            if self.preprocessor_speech:
+                data_src = self.preprocessor_speech(data_src, fs=self.fs)
+            speech, speech_lengths = extract_fbank(
+                data_src, data_type=self.data_type, frontend=self.frontend, is_final=True
+            )  # speech: [b, T, d]
+
+            if speech_lengths > self.batch_size:
+                continue
+            if self.permute:
+                speech = speech.permute(0, 2, 1)
+            asr_target = item["target"]
+            if self.preprocessor_text:
+                asr_target = self.preprocessor_text(asr_target)
+            emo_target = item.get("emo_target", "<|NEUTRAL|>")
+            event_target = item.get("event_target", "<|Speech|>")
+            text_language = item.get("text_language", "<|zh|>")
+            punc_itn_bottom = item.get("with_or_wo_itn", "<|woitn|>")
+
+            target_ids = self.tokenizer.encode(asr_target, allowed_special="all")
+            target_ids_len = len(target_ids)  # [text]
+            if target_ids_len > 200:
+                continue
+
+            lid_ids = self.tokenizer.encode(text_language, allowed_special="all")
+            emo_ids = self.tokenizer.encode(emo_target, allowed_special="all")
+            event_ids = self.tokenizer.encode(event_target, allowed_special="all")
+            punc_itn_bottom_ids = self.tokenizer.encode(punc_itn_bottom, allowed_special="all")
+
+            ids = lid_ids + emo_ids + event_ids + punc_itn_bottom_ids + target_ids # [lid, emo, lid, itn, text]
+            ids_lengths = len(ids)
+
+            text = torch.tensor(ids, dtype=torch.int64)
+            text_lengths = torch.tensor([ids_lengths], dtype=torch.int32)
+
+            output = {
+                "speech": speech[0, :, :],
+                "speech_lengths": speech_lengths,
+                "text": text,
+                "text_lengths": text_lengths,
+            }
+            break
+
+        return output
+
+    def collator(self, samples: list = None):
+        """Collator.
+        
+            Args:
+                samples: TODO.
+            """
+        outputs = {}
+        for sample in samples:
+            if sample is None:
+                continue
+            for key in sample.keys():
+                if key not in outputs:
+                    outputs[key] = []
+                outputs[key].append(sample[key])
+
+        if len(outputs) < 1:
+            logging.error(f"ERROR: data is empty!")
+            outputs = {
+                "speech": torch.rand((10, 128), dtype=torch.float32)[None, :, :],
+                "speech_lengths": torch.tensor(
+                    [
+                        10,
+                    ],
+                    dtype=torch.int32,
+                )[:, None],
+                "text": torch.tensor(
+                    [
+                        58836,
+                    ],
+                    dtype=torch.int32,
+                )[None, :],
+                "text_lengths": torch.tensor(
+                    [
+                        1,
+                    ],
+                    dtype=torch.int32,
+                )[:, None],
+            }
+            return outputs
+
+        for key, data_list in outputs.items():
+            if isinstance(data_list[0], torch.Tensor):
+                if data_list[0].dtype == torch.int64 or data_list[0].dtype == torch.int32:
+
+                    pad_value = self.int_pad_value
+                else:
+                    pad_value = self.float_pad_value
+
+                outputs[key] = torch.nn.utils.rnn.pad_sequence(
+                    data_list, batch_first=True, padding_value=pad_value
+                )
+
+        if self.batch_type != "example":
+            for i in range(10):
+                outputs = self._filter_badcase(outputs, i=i)
+
+        return outputs
+
+    def _filter_badcase(self, outputs, i=0):
+        """Internal: filter badcase.
+        
+            Args:
+                outputs: TODO.
+                i: TODO.
+            """
+        b, t, _ = outputs["speech"].shape
+
+        if b * t > self.batch_size * 1.25:
+            beg = torch.randint(0, 2, ()).item()
+            if b < 2:
+                beg = 0
+            logging.info(
+                f"Warning, b * t: {b * t} > {self.batch_size}, drop half data {i}th, beg:{beg}"
+            )
+            for key, data_list in outputs.items():
+                outputs[key] = outputs[key][beg : beg + b : 2]
+
+            speech_lengths_max = outputs["speech_lengths"].max().item()
+            outputs["speech"] = outputs["speech"][:, :speech_lengths_max, :]
+            text_lengths_max = outputs["text_lengths"].max().item()
+            outputs["text"] = outputs["text"][:, :text_lengths_max]
+
+        return outputs

funasr/download/download_dataset_from_hub.py

@@ -0,0 +1,20 @@
+def download_dataset():
+    """Download dataset."""
+    pass
+
+
+def download_dataset_from_ms(**kwargs):
+    """Download dataset from ms.
+    
+        Args:
+            **kwargs: Additional keyword arguments.
+        """
+    from modelscope.msdatasets import MsDataset
+
+    dataset_name = kwargs.get("dataset_name", "speech_asr/speech_asr_aishell1_trainsets")
+    subset_name = kwargs.get("subset_name", "default")
+    split = kwargs.get("split", "train")
+    data_dump_dir = kwargs.get("data_dump_dir", None)
+    ds = MsDataset.load(
+        dataset_name=dataset_name, subset_name=subset_name, split=split, cache_dir=data_dump_dir
+    )

funasr/download/download_model_from_hub.py

@@ -0,0 +1,292 @@
+import logging
+import os
+import json
+from omegaconf import OmegaConf, DictConfig
+
+from funasr.download.name_maps_from_hub import name_maps_ms, name_maps_hf, name_maps_openai
+
+
+def download_model(**kwargs):
+
+    """Download model from hub and parse its configuration.
+
+    Resolves model name aliases, downloads from ModelScope or HuggingFace,
+    reads config.yaml and configuration.json, and returns complete kwargs
+    for model instantiation.
+
+    Args:
+        **kwargs: Must include 'model' (str). Optional: 'hub', 'model_revision',
+            'is_training', etc.
+
+    Returns:
+        dict: Complete kwargs with resolved paths, model class name, and config.
+    """
+    hub = kwargs.get("hub", "ms")
+    if hub == "ms" or hub == "modelscope":
+        kwargs = download_from_ms(**kwargs)
+    elif hub == "hf" or hub == "huggingface":
+        kwargs = download_from_hf(**kwargs)
+    elif hub == "openai":
+        model_or_path = kwargs.get("model")
+        if os.path.exists(model_or_path):
+            # local path
+            kwargs["model_path"] = model_or_path
+            kwargs["model"] = "WhisperWarp"
+        else:
+            # model name
+            if model_or_path in name_maps_openai:
+                model_or_path = name_maps_openai[model_or_path]
+            kwargs["model_path"] = model_or_path
+
+    return kwargs
+
+
+def download_from_ms(**kwargs):
+    """Download from ms.
+    
+        Args:
+            **kwargs: Additional keyword arguments.
+        """
+    model_or_path = kwargs.get("model")
+    if model_or_path in name_maps_ms:
+        model_or_path = name_maps_ms[model_or_path]
+    model_revision = kwargs.get("model_revision", "master")
+    if not os.path.exists(model_or_path) and "model_path" not in kwargs:
+        try:
+            model_or_path = get_or_download_model_dir(
+                model_or_path,
+                model_revision,
+                is_training=kwargs.get("is_training"),
+                check_latest=kwargs.get("check_latest", True),
+            )
+        except Exception as e:
+            print(f"Download: {model_or_path} failed!: {e}")
+
+    kwargs["model_path"] = model_or_path if "model_path" not in kwargs else kwargs["model_path"]
+    model_or_path = kwargs["model_path"]
+    if os.path.exists(os.path.join(model_or_path, "configuration.json")):
+        with open(os.path.join(model_or_path, "configuration.json"), "r", encoding="utf-8") as f:
+            conf_json = json.load(f)
+
+            cfg = {}
+            if "file_path_metas" in conf_json:
+                add_file_root_path(model_or_path, conf_json["file_path_metas"], cfg)
+            # cfg.update(kwargs)
+            cfg = OmegaConf.merge(cfg, kwargs)
+            if "config" in cfg:
+                config = OmegaConf.load(cfg["config"])
+                kwargs = OmegaConf.merge(config, cfg)
+                kwargs["model"] = config["model"]
+    elif os.path.exists(os.path.join(model_or_path, "config.yaml")):
+        config = OmegaConf.load(os.path.join(model_or_path, "config.yaml"))
+        kwargs = OmegaConf.merge(config, kwargs)
+        init_param = os.path.join(model_or_path, "model.pt")
+        if "init_param" not in kwargs or not os.path.exists(kwargs["init_param"]):
+            kwargs["init_param"] = init_param
+            assert os.path.exists(kwargs["init_param"]), "init_param does not exist"
+        if os.path.exists(os.path.join(model_or_path, "tokens.txt")):
+            kwargs["tokenizer_conf"]["token_list"] = os.path.join(model_or_path, "tokens.txt")
+        if os.path.exists(os.path.join(model_or_path, "tokens.json")):
+            kwargs["tokenizer_conf"]["token_list"] = os.path.join(model_or_path, "tokens.json")
+        if os.path.exists(os.path.join(model_or_path, "seg_dict")):
+            kwargs["tokenizer_conf"]["seg_dict"] = os.path.join(model_or_path, "seg_dict")
+        if os.path.exists(os.path.join(model_or_path, "bpe.model")):
+            kwargs["tokenizer_conf"]["bpemodel"] = os.path.join(model_or_path, "bpe.model")
+        kwargs["model"] = config["model"]
+        if os.path.exists(os.path.join(model_or_path, "am.mvn")):
+            kwargs["frontend_conf"]["cmvn_file"] = os.path.join(model_or_path, "am.mvn")
+        if os.path.exists(os.path.join(model_or_path, "jieba_usr_dict")):
+            kwargs["jieba_usr_dict"] = os.path.join(model_or_path, "jieba_usr_dict")
+    if isinstance(kwargs, DictConfig):
+        kwargs = OmegaConf.to_container(kwargs, resolve=True)
+    logging.warning(f'trust_remote_code: {kwargs.get("trust_remote_code", False)}')
+    if os.path.exists(os.path.join(model_or_path, "requirements.txt")) and kwargs.get(
+        "trust_remote_code", False
+    ):
+        requirements = os.path.join(model_or_path, "requirements.txt")
+        print(f"Detect model requirements, begin to install it: {requirements}")
+        from funasr.utils.install_model_requirements import install_requirements
+
+        install_requirements(requirements)
+    if kwargs.get("trust_remote_code", False):
+        from funasr.utils.dynamic_import import import_module_from_path
+
+        model_code = kwargs.get("remote_code", "model")
+        import_module_from_path(model_code)
+
+        # from funasr.register import tables
+        # tables.print("model")
+    return kwargs
+
+
+def download_from_hf(**kwargs):
+    """Download from hf.
+    
+        Args:
+            **kwargs: Additional keyword arguments.
+        """
+    model_or_path = kwargs.get("model")
+    if model_or_path in name_maps_hf:
+        model_or_path = name_maps_hf[model_or_path]
+    model_revision = kwargs.get("model_revision", "master")
+    if not os.path.exists(model_or_path) and "model_path" not in kwargs:
+        try:
+            model_or_path = get_or_download_model_dir_hf(
+                model_or_path,
+                model_revision,
+                is_training=kwargs.get("is_training"),
+                check_latest=kwargs.get("check_latest", True),
+            )
+        except Exception as e:
+            print(f"Download: {model_or_path} failed!: {e}")
+
+    kwargs["model_path"] = model_or_path if "model_path" not in kwargs else kwargs["model_path"]
+
+    if os.path.exists(os.path.join(model_or_path, "configuration.json")):
+        with open(os.path.join(model_or_path, "configuration.json"), "r", encoding="utf-8") as f:
+            conf_json = json.load(f)
+
+            cfg = {}
+            if "file_path_metas" in conf_json:
+                add_file_root_path(model_or_path, conf_json["file_path_metas"], cfg)
+            cfg.update(kwargs)
+            if "config" in cfg:
+                config = OmegaConf.load(cfg["config"])
+                kwargs = OmegaConf.merge(config, cfg)
+                kwargs["model"] = config["model"]
+    elif os.path.exists(os.path.join(model_or_path, "config.yaml")) and os.path.exists(
+        os.path.join(model_or_path, "model.pt")
+    ):
+        config = OmegaConf.load(os.path.join(model_or_path, "config.yaml"))
+        kwargs = OmegaConf.merge(config, kwargs)
+        init_param = os.path.join(model_or_path, "model.pt")
+        kwargs["init_param"] = init_param
+        if os.path.exists(os.path.join(model_or_path, "tokens.txt")):
+            kwargs["tokenizer_conf"]["token_list"] = os.path.join(model_or_path, "tokens.txt")
+        if os.path.exists(os.path.join(model_or_path, "tokens.json")):
+            kwargs["tokenizer_conf"]["token_list"] = os.path.join(model_or_path, "tokens.json")
+        if os.path.exists(os.path.join(model_or_path, "seg_dict")):
+            kwargs["tokenizer_conf"]["seg_dict"] = os.path.join(model_or_path, "seg_dict")
+        if os.path.exists(os.path.join(model_or_path, "bpe.model")):
+            kwargs["tokenizer_conf"]["bpemodel"] = os.path.join(model_or_path, "bpe.model")
+        kwargs["model"] = config["model"]
+        if os.path.exists(os.path.join(model_or_path, "am.mvn")):
+            kwargs["frontend_conf"]["cmvn_file"] = os.path.join(model_or_path, "am.mvn")
+        if os.path.exists(os.path.join(model_or_path, "jieba_usr_dict")):
+            kwargs["jieba_usr_dict"] = os.path.join(model_or_path, "jieba_usr_dict")
+    if isinstance(kwargs, DictConfig):
+        kwargs = OmegaConf.to_container(kwargs, resolve=True)
+    logging.warning(f'trust_remote_code: {kwargs.get("trust_remote_code", False)}')
+    if os.path.exists(os.path.join(model_or_path, "requirements.txt")) and kwargs.get(
+        "trust_remote_code", False
+    ):
+        requirements = os.path.join(model_or_path, "requirements.txt")
+        print(f"Detect model requirements, begin to install it: {requirements}")
+        from funasr.utils.install_model_requirements import install_requirements
+
+        install_requirements(requirements)
+    return kwargs
+
+
+def add_file_root_path(model_or_path: str, file_path_metas: dict, cfg={}):
+
+    """Add file root path.
+    
+        Args:
+            model_or_path: TODO.
+            file_path_metas: TODO.
+            cfg: Configuration overrides.
+        """
+    if isinstance(file_path_metas, dict):
+        if isinstance(cfg, list):
+            cfg.append({})
+
+        for k, v in file_path_metas.items():
+            if isinstance(v, str):
+                p = os.path.join(model_or_path, v)
+                if os.path.exists(p):
+                    if isinstance(cfg, dict):
+                        cfg[k] = p
+                    elif isinstance(cfg, list):
+                        # if len(cfg) == 0:
+                        # cfg.append({})
+                        cfg[-1][k] = p
+
+            elif isinstance(v, dict):
+                if isinstance(cfg, dict):
+                    if k not in cfg:
+                        cfg[k] = {}
+                    add_file_root_path(model_or_path, v, cfg[k])
+                # elif isinstance(cfg, list):
+                #     cfg.append({})
+                #     add_file_root_path(model_or_path, v, cfg)
+            elif isinstance(v, (list, tuple)):
+                for i, vv in enumerate(v):
+                    if k not in cfg:
+                        cfg[k] = []
+                    if isinstance(vv, str):
+                        p = os.path.join(model_or_path, vv)
+                        # file_path_metas[i] = p
+                        if os.path.exists(p):
+                            if isinstance(cfg[k], dict):
+                                cfg[k] = p
+                            elif isinstance(cfg[k], list):
+                                cfg[k].append(p)
+                    elif isinstance(vv, dict):
+                        add_file_root_path(model_or_path, vv, cfg[k])
+
+    return cfg
+
+
+def get_or_download_model_dir(
+    model,
+    model_revision=None,
+    is_training=False,
+    check_latest=True,
+):
+    """Get local model directory or download model if necessary.
+
+    Args:
+        model (str): model id or path to local model directory.
+        model_revision  (str, optional): model version number.
+        :param is_training:
+    """
+    from modelscope.hub.check_model import check_local_model_is_latest
+    from modelscope.hub.snapshot_download import snapshot_download
+
+    from modelscope.utils.constant import Invoke, ThirdParty
+
+    key = Invoke.LOCAL_TRAINER if is_training else Invoke.PIPELINE
+
+    if os.path.exists(model) and check_latest:
+        model_cache_dir = model if os.path.isdir(model) else os.path.dirname(model)
+        try:
+            check_local_model_is_latest(
+                model_cache_dir, user_agent={Invoke.KEY: key, ThirdParty.KEY: "funasr"}
+            )
+        except:
+            print("could not check the latest version")
+    else:
+        model_cache_dir = snapshot_download(
+            model, revision=model_revision, user_agent={Invoke.KEY: key, ThirdParty.KEY: "funasr"}
+        )
+    return model_cache_dir
+
+
+def get_or_download_model_dir_hf(
+    model,
+    model_revision=None,
+    is_training=False,
+    check_latest=True,
+):
+    """Get local model directory or download model if necessary.
+
+    Args:
+        model (str): model id or path to local model directory.
+        model_revision  (str, optional): model version number.
+        :param is_training:
+    """
+    from huggingface_hub import snapshot_download
+
+    model_cache_dir = snapshot_download(model)
+    return model_cache_dir

funasr/download/file.py

@@ -0,0 +1,405 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import contextlib
+import os
+import tempfile
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+from typing import Generator, Union
+
+import requests
+from urllib.parse import urlparse
+
+
+def download_from_url(url):
+    """Download from url.
+    
+        Args:
+            url: TODO.
+        """
+    result = urlparse(url)
+    file_path = None
+    if result.scheme is not None and len(result.scheme) > 0:
+        storage = HTTPStorage()
+        # bytes
+        data = storage.read(url)
+        work_dir = tempfile.TemporaryDirectory().name
+        if not os.path.exists(work_dir):
+            os.makedirs(work_dir)
+        file_path = os.path.join(work_dir, os.path.basename(url))
+        with open(file_path, "wb") as fb:
+            fb.write(data)
+    assert file_path is not None, f"failed to download: {url}"
+    return file_path
+
+
+class Storage(metaclass=ABCMeta):
+    """Abstract class of storage.
+
+    All backends need to implement two apis: ``read()`` and ``read_text()``.
+    ``read()`` reads the file as a byte stream and ``read_text()`` reads
+    the file as texts.
+    """
+
+    @abstractmethod
+    def read(self, filepath: str):
+        """Read.
+        
+            Args:
+                filepath: TODO.
+            """
+        pass
+
+    @abstractmethod
+    def read_text(self, filepath: str):
+        """Read text.
+        
+            Args:
+                filepath: TODO.
+            """
+        pass
+
+    @abstractmethod
+    def write(self, obj: bytes, filepath: Union[str, Path]) -> None:
+        """Write.
+        
+            Args:
+                obj: TODO.
+                filepath: TODO.
+            """
+        pass
+
+    @abstractmethod
+    def write_text(self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8") -> None:
+        """Write text.
+        
+            Args:
+                obj: TODO.
+                filepath: TODO.
+                encoding: TODO.
+            """
+        pass
+
+
+class LocalStorage(Storage):
+    """Local hard disk storage"""
+
+    def read(self, filepath: Union[str, Path]) -> bytes:
+        """Read data from a given ``filepath`` with 'rb' mode.
+
+        Args:
+            filepath (str or Path): Path to read data.
+
+        Returns:
+            bytes: Expected bytes object.
+        """
+        with open(filepath, "rb") as f:
+            content = f.read()
+        return content
+
+    def read_text(self, filepath: Union[str, Path], encoding: str = "utf-8") -> str:
+        """Read data from a given ``filepath`` with 'r' mode.
+
+        Args:
+            filepath (str or Path): Path to read data.
+            encoding (str): The encoding format used to open the ``filepath``.
+                Default: 'utf-8'.
+
+        Returns:
+            str: Expected text reading from ``filepath``.
+        """
+        with open(filepath, "r", encoding=encoding) as f:
+            value_buf = f.read()
+        return value_buf
+
+    def write(self, obj: bytes, filepath: Union[str, Path]) -> None:
+        """Write data to a given ``filepath`` with 'wb' mode.
+
+        Note:
+            ``write`` will create a directory if the directory of ``filepath``
+            does not exist.
+
+        Args:
+            obj (bytes): Data to be written.
+            filepath (str or Path): Path to write data.
+        """
+        dirname = os.path.dirname(filepath)
+        if dirname and not os.path.exists(dirname):
+            os.makedirs(dirname, exist_ok=True)
+
+        with open(filepath, "wb") as f:
+            f.write(obj)
+
+    def write_text(self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8") -> None:
+        """Write data to a given ``filepath`` with 'w' mode.
+
+        Note:
+            ``write_text`` will create a directory if the directory of
+            ``filepath`` does not exist.
+
+        Args:
+            obj (str): Data to be written.
+            filepath (str or Path): Path to write data.
+            encoding (str): The encoding format used to open the ``filepath``.
+                Default: 'utf-8'.
+        """
+        dirname = os.path.dirname(filepath)
+        if dirname and not os.path.exists(dirname):
+            os.makedirs(dirname, exist_ok=True)
+
+        with open(filepath, "w", encoding=encoding) as f:
+            f.write(obj)
+
+    @contextlib.contextmanager
+    def as_local_path(self, filepath: Union[str, Path]) -> Generator[Union[str, Path], None, None]:
+        """Only for unified API and do nothing."""
+        yield filepath
+
+
+class HTTPStorage(Storage):
+    """HTTP and HTTPS storage."""
+
+    def read(self, url):
+        # TODO @wenmeng.zwm add progress bar if file is too large
+        """Read.
+        
+            Args:
+                url: TODO.
+            """
+        r = requests.get(url)
+        r.raise_for_status()
+        return r.content
+
+    def read_text(self, url):
+        """Read text.
+        
+            Args:
+                url: TODO.
+            """
+        r = requests.get(url)
+        r.raise_for_status()
+        return r.text
+
+    @contextlib.contextmanager
+    def as_local_path(self, filepath: str) -> Generator[Union[str, Path], None, None]:
+        """Download a file from ``filepath``.
+
+        ``as_local_path`` is decorated by :meth:`contextlib.contextmanager`. It
+        can be called with ``with`` statement, and when exists from the
+        ``with`` statement, the temporary path will be released.
+
+        Args:
+            filepath (str): Download a file from ``filepath``.
+
+        Examples:
+            >>> storage = HTTPStorage()
+            >>> # After existing from the ``with`` clause,
+            >>> # the path will be removed
+            >>> with storage.get_local_path('http://path/to/file') as path:
+            ...     # do something here
+        """
+        try:
+            f = tempfile.NamedTemporaryFile(delete=False)
+            f.write(self.read(filepath))
+            f.close()
+            yield f.name
+        finally:
+            os.remove(f.name)
+
+    def write(self, obj: bytes, url: Union[str, Path]) -> None:
+        """Write.
+        
+            Args:
+                obj: TODO.
+                url: TODO.
+            """
+        raise NotImplementedError("write is not supported by HTTP Storage")
+
+    def write_text(self, obj: str, url: Union[str, Path], encoding: str = "utf-8") -> None:
+        """Write text.
+        
+            Args:
+                obj: TODO.
+                url: TODO.
+                encoding: TODO.
+            """
+        raise NotImplementedError("write_text is not supported by HTTP Storage")
+
+
+class OSSStorage(Storage):
+    """OSS storage."""
+
+    def __init__(self, oss_config_file=None):
+        # read from config file or env var
+        """Initialize OSSStorage.
+        
+            Args:
+                oss_config_file: TODO.
+            """
+        raise NotImplementedError("OSSStorage.__init__ to be implemented in the future")
+
+    def read(self, filepath):
+        """Read.
+        
+            Args:
+                filepath: TODO.
+            """
+        raise NotImplementedError("OSSStorage.read to be implemented in the future")
+
+    def read_text(self, filepath, encoding="utf-8"):
+        """Read text.
+        
+            Args:
+                filepath: TODO.
+                encoding: TODO.
+            """
+        raise NotImplementedError("OSSStorage.read_text to be implemented in the future")
+
+    @contextlib.contextmanager
+    def as_local_path(self, filepath: str) -> Generator[Union[str, Path], None, None]:
+        """Download a file from ``filepath``.
+
+        ``as_local_path`` is decorated by :meth:`contextlib.contextmanager`. It
+        can be called with ``with`` statement, and when exists from the
+        ``with`` statement, the temporary path will be released.
+
+        Args:
+            filepath (str): Download a file from ``filepath``.
+
+        Examples:
+            >>> storage = OSSStorage()
+            >>> # After existing from the ``with`` clause,
+            >>> # the path will be removed
+            >>> with storage.get_local_path('http://path/to/file') as path:
+            ...     # do something here
+        """
+        try:
+            f = tempfile.NamedTemporaryFile(delete=False)
+            f.write(self.read(filepath))
+            f.close()
+            yield f.name
+        finally:
+            os.remove(f.name)
+
+    def write(self, obj: bytes, filepath: Union[str, Path]) -> None:
+        """Write.
+        
+            Args:
+                obj: TODO.
+                filepath: TODO.
+            """
+        raise NotImplementedError("OSSStorage.write to be implemented in the future")
+
+    def write_text(self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8") -> None:
+        """Write text.
+        
+            Args:
+                obj: TODO.
+                filepath: TODO.
+                encoding: TODO.
+            """
+        raise NotImplementedError("OSSStorage.write_text to be implemented in the future")
+
+
+G_STORAGES = {}
+
+
+class File(object):
+    _prefix_to_storage: dict = {
+        "oss": OSSStorage,
+        "http": HTTPStorage,
+        "https": HTTPStorage,
+        "local": LocalStorage,
+    }
+
+    @staticmethod
+    def _get_storage(uri):
+        """Internal: get storage.
+        
+            Args:
+                uri: TODO.
+            """
+        assert isinstance(uri, str), f"uri should be str type, but got {type(uri)}"
+
+        if "://" not in uri:
+            # local path
+            storage_type = "local"
+        else:
+            prefix, _ = uri.split("://")
+            storage_type = prefix
+
+        assert storage_type in File._prefix_to_storage, (
+            f"Unsupported uri {uri}, valid prefixs: " f"{list(File._prefix_to_storage.keys())}"
+        )
+
+        if storage_type not in G_STORAGES:
+            G_STORAGES[storage_type] = File._prefix_to_storage[storage_type]()
+
+        return G_STORAGES[storage_type]
+
+    @staticmethod
+    def read(uri: str) -> bytes:
+        """Read data from a given ``filepath`` with 'rb' mode.
+
+        Args:
+            filepath (str or Path): Path to read data.
+
+        Returns:
+            bytes: Expected bytes object.
+        """
+        storage = File._get_storage(uri)
+        return storage.read(uri)
+
+    @staticmethod
+    def read_text(uri: Union[str, Path], encoding: str = "utf-8") -> str:
+        """Read data from a given ``filepath`` with 'r' mode.
+
+        Args:
+            filepath (str or Path): Path to read data.
+            encoding (str): The encoding format used to open the ``filepath``.
+                Default: 'utf-8'.
+
+        Returns:
+            str: Expected text reading from ``filepath``.
+        """
+        storage = File._get_storage(uri)
+        return storage.read_text(uri)
+
+    @staticmethod
+    def write(obj: bytes, uri: Union[str, Path]) -> None:
+        """Write data to a given ``filepath`` with 'wb' mode.
+
+        Note:
+            ``write`` will create a directory if the directory of ``filepath``
+            does not exist.
+
+        Args:
+            obj (bytes): Data to be written.
+            filepath (str or Path): Path to write data.
+        """
+        storage = File._get_storage(uri)
+        return storage.write(obj, uri)
+
+    @staticmethod
+    def write_text(obj: str, uri: str, encoding: str = "utf-8") -> None:
+        """Write data to a given ``filepath`` with 'w' mode.
+
+        Note:
+            ``write_text`` will create a directory if the directory of
+            ``filepath`` does not exist.
+
+        Args:
+            obj (str): Data to be written.
+            filepath (str or Path): Path to write data.
+            encoding (str): The encoding format used to open the ``filepath``.
+                Default: 'utf-8'.
+        """
+        storage = File._get_storage(uri)
+        return storage.write_text(obj, uri)
+
+    @contextlib.contextmanager
+    def as_local_path(uri: str) -> Generator[Union[str, Path], None, None]:
+        """Only for unified API and do nothing."""
+        storage = File._get_storage(uri)
+        with storage.as_local_path(uri) as local_path:
+            yield local_path

funasr/download/name_maps_from_hub.py

@@ -0,0 +1,57 @@
+name_maps_ms = {
+    "paraformer": "iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch",
+    "paraformer-zh": "iic/speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch",
+    "paraformer-en": "iic/speech_paraformer-large-vad-punc_asr_nat-en-16k-common-vocab10020",
+    "paraformer-en-spk": "iic/speech_paraformer-large-vad-punc_asr_nat-en-16k-common-vocab10020",
+    "paraformer-zh-streaming": "iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-online",
+    "fsmn-vad": "iic/speech_fsmn_vad_zh-cn-16k-common-pytorch",
+    "ct-punc": "iic/punc_ct-transformer_cn-en-common-vocab471067-large",
+    "ct-punc-c": "iic/punc_ct-transformer_zh-cn-common-vocab272727-pytorch",
+    "fa-zh": "iic/speech_timestamp_prediction-v1-16k-offline",
+    "cam++": "iic/speech_campplus_sv_zh-cn_16k-common",
+    "Whisper-large-v2": "iic/speech_whisper-large_asr_multilingual",
+    "Whisper-large-v3": "iic/Whisper-large-v3",
+    "Qwen-Audio": "Qwen/Qwen-Audio",
+    "emotion2vec_plus_large": "iic/emotion2vec_plus_large",
+    "emotion2vec_plus_base": "iic/emotion2vec_plus_base",
+    "emotion2vec_plus_seed": "iic/emotion2vec_plus_seed",
+    "Whisper-large-v3-turbo": "iic/Whisper-large-v3-turbo",
+}
+
+name_maps_hf = {
+    "paraformer": "funasr/paraformer-zh",
+    "paraformer-zh": "funasr/paraformer-zh",
+    "paraformer-en": "funasr/paraformer-zh",
+    "paraformer-zh-streaming": "funasr/paraformer-zh-streaming",
+    "fsmn-vad": "funasr/fsmn-vad",
+    "ct-punc": "funasr/ct-punc",
+    "ct-punc-c": "iic/punc_ct-transformer_zh-cn-common-vocab272727-pytorch",
+    "fa-zh": "funasr/fa-zh",
+    "cam++": "funasr/campplus",
+    "Whisper-large-v2": "iic/speech_whisper-large_asr_multilingual",
+    "Whisper-large-v3": "iic/Whisper-large-v3",
+    "Qwen-Audio": "Qwen/Qwen-Audio",
+    "emotion2vec_plus_large": "emotion2vec/emotion2vec_plus_large",
+    "iic/emotion2vec_plus_large": "emotion2vec/emotion2vec_plus_large",
+    "emotion2vec_plus_base": "emotion2vec/emotion2vec_plus_base",
+    "iic/emotion2vec_plus_base": "emotion2vec/emotion2vec_plus_base",
+    "emotion2vec_plus_seed": "emotion2vec/emotion2vec_plus_seed",
+    "iic/emotion2vec_plus_seed": "emotion2vec/emotion2vec_plus_seed",
+    "Whisper-large-v3-turbo": "iic/Whisper-large-v3-turbo",
+}
+
+name_maps_openai = {
+    "Whisper-tiny.en": "tiny.en",
+    "Whisper-tiny": "tiny",
+    "Whisper-base.en": "base.en",
+    "Whisper-base": "base",
+    "Whisper-small.en": "small.en",
+    "Whisper-small": "small",
+    "Whisper-medium.en": "medium.en",
+    "Whisper-medium": "medium",
+    "Whisper-large-v1": "large-v1",
+    "Whisper-large-v2": "large-v2",
+    "Whisper-large-v3": "large-v3",
+    "Whisper-large": "large",
+    "Whisper-large-v3-turbo": "turbo",
+}

funasr/download/runtime_sdk_download_tool.py

@@ -0,0 +1,59 @@
+import os
+import argparse
+from pathlib import Path
+
+from funasr.utils.type_utils import str2bool
+
+
+def main():
+    """Main."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-name", type=str, required=True)
+    parser.add_argument("--export-dir", type=str, required=True)
+    parser.add_argument("--export", type=str2bool, default=True, help="whether to export model")
+    parser.add_argument("--type", type=str, default="onnx", help='["onnx", "torchscript", "bladedisc"]')
+    parser.add_argument("--device", type=str, default="cpu", help='["cpu", "cuda"]')
+    parser.add_argument("--quantize", type=str2bool, default=False, help="export quantized model")
+    parser.add_argument("--fallback-num", type=int, default=0, help="amp fallback number")
+    parser.add_argument("--audio_in", type=str, default=None, help='["wav", "wav.scp"]')
+    parser.add_argument("--model_revision", type=str, default=None, help="model_revision")
+    parser.add_argument("--calib_num", type=int, default=200, help="calib max num")
+    args = parser.parse_args()
+
+    model_dir = args.model_name
+    output_dir = args.model_name
+    if not Path(args.model_name).exists():
+        from modelscope.hub.snapshot_download import snapshot_download
+
+        try:
+            model_dir = snapshot_download(
+                args.model_name, cache_dir=args.export_dir, revision=args.model_revision
+            )
+            output_dir = os.path.join(args.export_dir, args.model_name)
+        except:
+            raise "model_dir must be model_name in modelscope or local path downloaded from modelscope, but is {}".format(
+                model_dir
+            )
+    if args.export:
+        model_file = os.path.join(model_dir, "model.onnx")
+        if args.quantize:
+            model_file = os.path.join(model_dir, "model_quant.onnx")
+        if args.type == "torchscript":
+            model_file = os.path.join(model_dir, "model.torchscript")
+            args.device = "cuda"
+        elif args.type == "bladedisc":
+            model_file = os.path.join(model_dir, "model_blade.torchscript")
+            args.device = "cuda"
+        if not os.path.exists(model_file):
+            print("model is not exist, begin to export " + model_file)
+            from funasr import AutoModel
+
+            export_model = AutoModel(model=args.model_name, output_dir=output_dir, device=args.device)
+            export_model.export(
+                    quantize=args.quantize,
+                    type=args.type,
+                    )
+
+
+if __name__ == "__main__":
+    main()

funasr/frontends/abs_frontend.py

@@ -0,0 +1,12 @@
+from abc import ABC, abstractmethod
+import torch
+
+
+class AbsFrontend(ABC, torch.nn.Module):
+    @abstractmethod
+    def output_size(self) -> int:
+        raise NotImplementedError
+
+    @abstractmethod
+    def forward(self, input, input_lengths):
+        raise NotImplementedError

funasr/frontends/default.py

@@ -0,0 +1,416 @@
+import copy
+from typing import Optional
+from typing import Tuple
+from typing import Union
+import logging
+import numpy as np
+import torch
+import torch.nn as nn
+
+try:
+    from torch_complex.tensor import ComplexTensor
+except:
+    print("Please install torch_complex firstly")
+
+from funasr.frontends.utils.log_mel import LogMel
+from funasr.frontends.utils.stft import Stft
+from funasr.frontends.utils.frontend import Frontend
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+from funasr.register import tables
+
+
+@tables.register("frontend_classes", "DefaultFrontend")
+@tables.register("frontend_classes", "EspnetFrontend")
+class DefaultFrontend(nn.Module):
+    """Conventional frontend structure for ASR.
+    Stft -> WPE -> MVDR-Beamformer -> Power-spec -> Mel-Fbank -> CMVN
+    """
+
+    def __init__(
+        self,
+        fs: int = 16000,
+        n_fft: int = 512,
+        win_length: int = None,
+        hop_length: int = 128,
+        window: Optional[str] = "hann",
+        center: bool = True,
+        normalized: bool = False,
+        onesided: bool = True,
+        n_mels: int = 80,
+        fmin: int = None,
+        fmax: int = None,
+        htk: bool = False,
+        frontend_conf: Optional[dict] = None,
+        apply_stft: bool = True,
+        use_channel: int = None,
+        **kwargs,
+    ):
+        """Initialize DefaultFrontend.
+        
+            Args:
+                fs: TODO.
+                n_fft: TODO.
+                win_length: TODO.
+                hop_length: TODO.
+                window: TODO.
+                center: TODO.
+                normalized: TODO.
+                onesided: TODO.
+                n_mels: TODO.
+                fmin: TODO.
+                fmax: TODO.
+                htk: TODO.
+                frontend_conf: Configuration dict for frontend.
+                apply_stft: TODO.
+                use_channel: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+
+        # Deepcopy (In general, dict shouldn't be used as default arg)
+        frontend_conf = copy.deepcopy(frontend_conf)
+        self.hop_length = hop_length
+        self.fs = fs
+
+        if apply_stft:
+            self.stft = Stft(
+                n_fft=n_fft,
+                win_length=win_length,
+                hop_length=hop_length,
+                center=center,
+                window=window,
+                normalized=normalized,
+                onesided=onesided,
+            )
+        else:
+            self.stft = None
+        self.apply_stft = apply_stft
+
+        if frontend_conf is not None:
+            self.frontend = Frontend(idim=n_fft // 2 + 1, **frontend_conf)
+        else:
+            self.frontend = None
+
+        self.logmel = LogMel(
+            fs=fs,
+            n_fft=n_fft,
+            n_mels=n_mels,
+            fmin=fmin,
+            fmax=fmax,
+            htk=htk,
+        )
+        self.n_mels = n_mels
+        self.use_channel = use_channel
+        self.frontend_type = "default"
+
+    def output_size(self) -> int:
+        """Output size."""
+        return self.n_mels
+
+    def forward(
+        self, input: torch.Tensor, input_lengths: Union[torch.Tensor, list]
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass for training.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+            """
+        if isinstance(input_lengths, list):
+            input_lengths = torch.tensor(input_lengths)
+        if input.dtype == torch.float64:
+            input = input.float()
+        # 1. Domain-conversion: e.g. Stft: time -> time-freq
+        if self.stft is not None:
+            input_stft, feats_lens = self._compute_stft(input, input_lengths)
+        else:
+            input_stft = ComplexTensor(input[..., 0], input[..., 1])
+            feats_lens = input_lengths
+        # 2. [Option] Speech enhancement
+        if self.frontend is not None:
+            assert isinstance(input_stft, ComplexTensor), type(input_stft)
+            # input_stft: (Batch, Length, [Channel], Freq)
+            input_stft, _, mask = self.frontend(input_stft, feats_lens)
+
+        # 3. [Multi channel case]: Select a channel
+        if input_stft.dim() == 4:
+            # h: (B, T, C, F) -> h: (B, T, F)
+            if self.training:
+                if self.use_channel is not None:
+                    input_stft = input_stft[:, :, self.use_channel, :]
+                else:
+                    # Select 1ch randomly
+                    ch = np.random.randint(input_stft.size(2))
+                    input_stft = input_stft[:, :, ch, :]
+            else:
+                # Use the first channel
+                input_stft = input_stft[:, :, 0, :]
+
+        # 4. STFT -> Power spectrum
+        # h: ComplexTensor(B, T, F) -> torch.Tensor(B, T, F)
+        input_power = input_stft.real**2 + input_stft.imag**2
+
+        # 5. Feature transform e.g. Stft -> Log-Mel-Fbank
+        # input_power: (Batch, [Channel,] Length, Freq)
+        #       -> input_feats: (Batch, Length, Dim)
+        input_feats, _ = self.logmel(input_power, feats_lens)
+
+        return input_feats, feats_lens
+
+    def _compute_stft(self, input: torch.Tensor, input_lengths: torch.Tensor) -> torch.Tensor:
+        """Internal: compute stft.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+            """
+        input_stft, feats_lens = self.stft(input, input_lengths)
+
+        assert input_stft.dim() >= 4, input_stft.shape
+        # "2" refers to the real/imag parts of Complex
+        assert input_stft.shape[-1] == 2, input_stft.shape
+
+        # Change torch.Tensor to ComplexTensor
+        # input_stft: (..., F, 2) -> (..., F)
+        input_stft = ComplexTensor(input_stft[..., 0], input_stft[..., 1])
+        return input_stft, feats_lens
+
+
+class MultiChannelFrontend(nn.Module):
+    """Conventional frontend structure for ASR.
+    Stft -> WPE -> MVDR-Beamformer -> Power-spec -> Mel-Fbank -> CMVN
+    """
+
+    def __init__(
+        self,
+        fs: int = 16000,
+        n_fft: int = 512,
+        win_length: int = None,
+        hop_length: int = None,
+        frame_length: int = None,
+        frame_shift: int = None,
+        window: Optional[str] = "hann",
+        center: bool = True,
+        normalized: bool = False,
+        onesided: bool = True,
+        n_mels: int = 80,
+        fmin: int = None,
+        fmax: int = None,
+        htk: bool = False,
+        frontend_conf: Optional[dict] = None,
+        apply_stft: bool = True,
+        use_channel: int = None,
+        lfr_m: int = 1,
+        lfr_n: int = 1,
+        cmvn_file: str = None,
+        mc: bool = True,
+    ):
+        """Initialize MultiChannelFrontend.
+        
+            Args:
+                fs: TODO.
+                n_fft: TODO.
+                win_length: TODO.
+                hop_length: TODO.
+                frame_length: TODO.
+                frame_shift: TODO.
+                window: TODO.
+                center: TODO.
+                normalized: TODO.
+                onesided: TODO.
+                n_mels: TODO.
+                fmin: TODO.
+                fmax: TODO.
+                htk: TODO.
+                frontend_conf: Configuration dict for frontend.
+                apply_stft: TODO.
+                use_channel: TODO.
+                lfr_m: TODO.
+                lfr_n: TODO.
+                cmvn_file: TODO.
+                mc: TODO.
+            """
+        super().__init__()
+        # Deepcopy (In general, dict shouldn't be used as default arg)
+        frontend_conf = copy.deepcopy(frontend_conf)
+        if win_length is None and hop_length is None:
+            self.win_length = frame_length * 16
+            self.hop_length = frame_shift * 16
+        elif frame_length is None and frame_shift is None:
+            self.win_length = self.win_length
+            self.hop_length = self.hop_length
+        else:
+            logging.error(
+                "Only one of (win_length, hop_length) and (frame_length, frame_shift)" "can be set."
+            )
+            exit(1)
+
+        if apply_stft:
+            self.stft = Stft(
+                n_fft=n_fft,
+                win_length=self.win_length,
+                hop_length=self.hop_length,
+                center=center,
+                window=window,
+                normalized=normalized,
+                onesided=onesided,
+            )
+        else:
+            self.stft = None
+        self.apply_stft = apply_stft
+
+        if frontend_conf is not None:
+            self.frontend = Frontend(idim=n_fft // 2 + 1, **frontend_conf)
+        else:
+            self.frontend = None
+
+        self.logmel = LogMel(
+            fs=fs,
+            n_fft=n_fft,
+            n_mels=n_mels,
+            fmin=fmin,
+            fmax=fmax,
+            htk=htk,
+        )
+        self.n_mels = n_mels
+        self.use_channel = use_channel
+        self.mc = mc
+        if not self.mc:
+            if self.use_channel is not None:
+                logging.info("use the channel %d" % (self.use_channel))
+            else:
+                logging.info("random select channel")
+            self.cmvn_file = cmvn_file
+            if self.cmvn_file is not None:
+                mean, std = self._load_cmvn(self.cmvn_file)
+                self.register_buffer("mean", torch.from_numpy(mean))
+                self.register_buffer("std", torch.from_numpy(std))
+        self.frontend_type = "multichannelfrontend"
+
+    def output_size(self) -> int:
+        """Output size."""
+        return self.n_mels
+
+    def forward(
+        self, input: torch.Tensor, input_lengths: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # 1. Domain-conversion: e.g. Stft: time -> time-freq
+        """Forward pass for training.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+            """
+        if self.stft is not None:
+            input_stft, feats_lens = self._compute_stft(input, input_lengths)
+        else:
+            input_stft = ComplexTensor(input[..., 0], input[..., 1])
+            feats_lens = input_lengths
+        # 2. [Option] Speech enhancement
+        if self.frontend is not None:
+            assert isinstance(input_stft, ComplexTensor), type(input_stft)
+            # input_stft: (Batch, Length, [Channel], Freq)
+            input_stft, _, mask = self.frontend(input_stft, feats_lens)
+
+        # 3. [Multi channel case]: Select a channel(sa_asr)
+        if input_stft.dim() == 4 and not self.mc:
+            # h: (B, T, C, F) -> h: (B, T, F)
+            if self.training:
+                if self.use_channel is not None:
+                    input_stft = input_stft[:, :, self.use_channel, :]
+
+                else:
+                    # Select 1ch randomly
+                    ch = np.random.randint(input_stft.size(2))
+                    input_stft = input_stft[:, :, ch, :]
+            else:
+                # Use the first channel
+                input_stft = input_stft[:, :, 0, :]
+
+        # 4. STFT -> Power spectrum
+        # h: ComplexTensor(B, T, F) -> torch.Tensor(B, T, F)
+        input_power = input_stft.real**2 + input_stft.imag**2
+
+        # 5. Feature transform e.g. Stft -> Log-Mel-Fbank
+        # input_power: (Batch, [Channel,] Length, Freq)
+        #       -> input_feats: (Batch, Length, Dim)
+        input_feats, _ = self.logmel(input_power, feats_lens)
+        if self.mc:
+            # MFCCA
+            if input_feats.dim() == 4:
+                bt = input_feats.size(0)
+                channel_size = input_feats.size(2)
+                input_feats = (
+                    input_feats.transpose(1, 2).reshape(bt * channel_size, -1, 80).contiguous()
+                )
+                feats_lens = feats_lens.repeat(1, channel_size).squeeze()
+            else:
+                channel_size = 1
+            return input_feats, feats_lens, channel_size
+        else:
+            # 6. Apply CMVN
+            if self.cmvn_file is not None:
+                if feats_lens is None:
+                    feats_lens = input_feats.new_full([input_feats.size(0)], input_feats.size(1))
+                self.mean = self.mean.to(input_feats.device, input_feats.dtype)
+                self.std = self.std.to(input_feats.device, input_feats.dtype)
+                mask = make_pad_mask(feats_lens, input_feats, 1)
+
+                if input_feats.requires_grad:
+                    input_feats = input_feats + self.mean
+                else:
+                    input_feats += self.mean
+                if input_feats.requires_grad:
+                    input_feats = input_feats.masked_fill(mask, 0.0)
+                else:
+                    input_feats.masked_fill_(mask, 0.0)
+
+                input_feats *= self.std
+
+            return input_feats, feats_lens
+
+    def _compute_stft(self, input: torch.Tensor, input_lengths: torch.Tensor) -> torch.Tensor:
+        """Internal: compute stft.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+            """
+        input_stft, feats_lens = self.stft(input, input_lengths)
+
+        assert input_stft.dim() >= 4, input_stft.shape
+        # "2" refers to the real/imag parts of Complex
+        assert input_stft.shape[-1] == 2, input_stft.shape
+
+        # Change torch.Tensor to ComplexTensor
+        # input_stft: (..., F, 2) -> (..., F)
+        input_stft = ComplexTensor(input_stft[..., 0], input_stft[..., 1])
+        return input_stft, feats_lens
+
+    def _load_cmvn(self, cmvn_file):
+        """Internal: load cmvn.
+        
+            Args:
+                cmvn_file: TODO.
+            """
+        with open(cmvn_file, "r", encoding="utf-8") as f:
+            lines = f.readlines()
+        means_list = []
+        vars_list = []
+        for i in range(len(lines)):
+            line_item = lines[i].split()
+            if line_item[0] == "<AddShift>":
+                line_item = lines[i + 1].split()
+                if line_item[0] == "<LearnRateCoef>":
+                    add_shift_line = line_item[3 : (len(line_item) - 1)]
+                    means_list = list(add_shift_line)
+                    continue
+            elif line_item[0] == "<Rescale>":
+                line_item = lines[i + 1].split()
+                if line_item[0] == "<LearnRateCoef>":
+                    rescale_line = line_item[3 : (len(line_item) - 1)]
+                    vars_list = list(rescale_line)
+                    continue
+        means = np.array(means_list).astype(np.float)
+        vars = np.array(vars_list).astype(np.float)
+        return means, vars

funasr/frontends/eend_ola_feature.py

@@ -0,0 +1,73 @@
+# Copyright 2019 Hitachi, Ltd. (author: Yusuke Fujita)
+# Licensed under the MIT license.
+#
+# This module is for computing audio features
+
+import librosa
+import numpy as np
+
+
+def transform(Y, dtype=np.float32):
+    """Transform.
+    
+        Args:
+            Y: TODO.
+            dtype: TODO.
+        """
+    Y = np.abs(Y)
+    n_fft = 2 * (Y.shape[1] - 1)
+    sr = 8000
+    n_mels = 23
+    mel_basis = librosa.filters.mel(sr, n_fft, n_mels)
+    Y = np.dot(Y**2, mel_basis.T)
+    Y = np.log10(np.maximum(Y, 1e-10))
+    mean = np.mean(Y, axis=0)
+    Y = Y - mean
+    return Y.astype(dtype)
+
+
+def subsample(Y, T, subsampling=1):
+    """Subsample.
+    
+        Args:
+            Y: TODO.
+            T: TODO.
+            subsampling: TODO.
+        """
+    Y_ss = Y[::subsampling]
+    T_ss = T[::subsampling]
+    return Y_ss, T_ss
+
+
+def splice(Y, context_size=0):
+    """Splice.
+    
+        Args:
+            Y: TODO.
+            context_size: Size/dimension parameter.
+        """
+    Y_pad = np.pad(Y, [(context_size, context_size), (0, 0)], "constant")
+    Y_spliced = np.lib.stride_tricks.as_strided(
+        np.ascontiguousarray(Y_pad),
+        (Y.shape[0], Y.shape[1] * (2 * context_size + 1)),
+        (Y.itemsize * Y.shape[1], Y.itemsize),
+        writeable=False,
+    )
+    return Y_spliced
+
+
+def stft(data, frame_size=1024, frame_shift=256):
+    """Stft.
+    
+        Args:
+            data: TODO.
+            frame_size: Size/dimension parameter.
+            frame_shift: TODO.
+        """
+    fft_size = 1 << (frame_size - 1).bit_length()
+    if len(data) % frame_shift == 0:
+        return librosa.stft(data, n_fft=fft_size, win_length=frame_size, hop_length=frame_shift).T[
+            :-1
+        ]
+    else:
+        return librosa.stft(data, n_fft=fft_size, win_length=frame_size, hop_length=frame_shift).T

funasr/frontends/fused.py

@@ -0,0 +1,157 @@
+from funasr.frontends.default import DefaultFrontend
+from funasr.frontends.s3prl import S3prlFrontend
+import numpy as np
+import torch
+import torch.nn as nn
+from typing import Tuple
+
+
+class FusedFrontends(nn.Module):
+    def __init__(self, frontends=None, align_method="linear_projection", proj_dim=100, fs=16000):
+
+        """Initialize FusedFrontends.
+        
+            Args:
+                frontends: TODO.
+                align_method: TODO.
+                proj_dim: Size/dimension parameter.
+                fs: TODO.
+            """
+        super().__init__()
+        self.align_method = align_method  # fusing method : linear_projection only for now
+        self.proj_dim = proj_dim  # dim of the projection done on each frontend
+        self.frontends = []  # list of the frontends to combine
+
+        for i, frontend in enumerate(frontends):
+            frontend_type = frontend["frontend_type"]
+            if frontend_type == "default":
+                n_mels, fs, n_fft, win_length, hop_length = (
+                    frontend.get("n_mels", 80),
+                    fs,
+                    frontend.get("n_fft", 512),
+                    frontend.get("win_length"),
+                    frontend.get("hop_length", 128),
+                )
+                window, center, normalized, onesided = (
+                    frontend.get("window", "hann"),
+                    frontend.get("center", True),
+                    frontend.get("normalized", False),
+                    frontend.get("onesided", True),
+                )
+                fmin, fmax, htk, apply_stft = (
+                    frontend.get("fmin", None),
+                    frontend.get("fmax", None),
+                    frontend.get("htk", False),
+                    frontend.get("apply_stft", True),
+                )
+
+                self.frontends.append(
+                    DefaultFrontend(
+                        n_mels=n_mels,
+                        n_fft=n_fft,
+                        fs=fs,
+                        win_length=win_length,
+                        hop_length=hop_length,
+                        window=window,
+                        center=center,
+                        normalized=normalized,
+                        onesided=onesided,
+                        fmin=fmin,
+                        fmax=fmax,
+                        htk=htk,
+                        apply_stft=apply_stft,
+                    )
+                )
+            elif frontend_type == "s3prl":
+                frontend_conf, download_dir, multilayer_feature = (
+                    frontend.get("frontend_conf"),
+                    frontend.get("download_dir"),
+                    frontend.get("multilayer_feature"),
+                )
+                self.frontends.append(
+                    S3prlFrontend(
+                        fs=fs,
+                        frontend_conf=frontend_conf,
+                        download_dir=download_dir,
+                        multilayer_feature=multilayer_feature,
+                    )
+                )
+
+            else:
+                raise NotImplementedError  # frontends are only default or s3prl
+
+        self.frontends = torch.nn.ModuleList(self.frontends)
+
+        self.gcd = np.gcd.reduce([frontend.hop_length for frontend in self.frontends])
+        self.factors = [frontend.hop_length // self.gcd for frontend in self.frontends]
+        if torch.cuda.is_available():
+            dev = "cuda"
+        elif torch.xpu.is_available():
+            dev = "xpu"
+        elif torch.backends.mps.is_available():
+            dev = "mps"
+        else:
+            dev = "cpu"
+        if self.align_method == "linear_projection":
+            self.projection_layers = [
+                torch.nn.Linear(
+                    in_features=frontend.output_size(),
+                    out_features=self.factors[i] * self.proj_dim,
+                )
+                for i, frontend in enumerate(self.frontends)
+            ]
+            self.projection_layers = torch.nn.ModuleList(self.projection_layers)
+            self.projection_layers = self.projection_layers.to(torch.device(dev))
+
+    def output_size(self) -> int:
+        """Output size."""
+        return len(self.frontends) * self.proj_dim
+
+    def forward(
+        self, input: torch.Tensor, input_lengths: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+
+        # step 0 : get all frontends features
+        """Forward pass for training.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+            """
+        self.feats = []
+        for frontend in self.frontends:
+            with torch.no_grad():
+                input_feats, feats_lens = frontend.forward(input, input_lengths)
+            self.feats.append([input_feats, feats_lens])
+
+        if self.align_method == "linear_projection":  # TODO(Dan): to add other align methods
+
+            # first step : projections
+            self.feats_proj = []
+            for i, frontend in enumerate(self.frontends):
+                input_feats = self.feats[i][0]
+                self.feats_proj.append(self.projection_layers[i](input_feats))
+
+            # 2nd step : reshape
+            self.feats_reshaped = []
+            for i, frontend in enumerate(self.frontends):
+                input_feats_proj = self.feats_proj[i]
+                bs, nf, dim = input_feats_proj.shape
+                input_feats_reshaped = torch.reshape(
+                    input_feats_proj, (bs, nf * self.factors[i], dim // self.factors[i])
+                )
+                self.feats_reshaped.append(input_feats_reshaped)
+
+            # 3rd step : drop the few last frames
+            m = min([x.shape[1] for x in self.feats_reshaped])
+            self.feats_final = [x[:, :m, :] for x in self.feats_reshaped]
+
+            input_feats = torch.cat(
+                self.feats_final, dim=-1
+            )  # change the input size of the preencoder : proj_dim * n_frontends
+            feats_lens = torch.ones_like(self.feats[0][1]) * (m)
+
+        else:
+            raise NotImplementedError
+
+        return input_feats, feats_lens

funasr/frontends/s3prl.py

@@ -0,0 +1,166 @@
+import copy
+import logging
+import os
+from argparse import Namespace
+from typing import Optional
+from typing import Tuple
+from typing import Union
+
+try:
+    import humanfriendly
+except ImportError:
+    humanfriendly = None
+import torch
+import torch.nn as nn
+
+from funasr.frontends.utils.frontend import Frontend
+from funasr.models.transformer.utils.nets_utils import pad_list
+
+
+def base_s3prl_setup(args):
+    """Base s3prl setup.
+    
+        Args:
+            args: TODO.
+        """
+    args.upstream_feature_selection = getattr(args, "upstream_feature_selection", None)
+    args.upstream_model_config = getattr(args, "upstream_model_config", None)
+    args.upstream_refresh = getattr(args, "upstream_refresh", False)
+    args.upstream_ckpt = getattr(args, "upstream_ckpt", None)
+    args.init_ckpt = getattr(args, "init_ckpt", None)
+    args.verbose = getattr(args, "verbose", False)
+    args.tile_factor = getattr(args, "tile_factor", 1)
+    return args
+
+
+class S3prlFrontend(nn.Module):
+    """Speech Pretrained Representation frontend structure for ASR."""
+
+    def __init__(
+        self,
+        fs: Union[int, str] = 16000,
+        frontend_conf: Optional[dict] = None,
+        download_dir: str = None,
+        multilayer_feature: bool = False,
+    ):
+        """Initialize S3prlFrontend.
+        
+            Args:
+                fs: TODO.
+                frontend_conf: Configuration dict for frontend.
+                download_dir: TODO.
+                multilayer_feature: TODO.
+            """
+        super().__init__()
+        if isinstance(fs, str):
+            if humanfriendly is not None:
+                fs = humanfriendly.parse_size(fs)
+            else:
+                fs = int(fs)
+
+        if download_dir is not None:
+            torch.hub.set_dir(download_dir)
+
+        self.multilayer_feature = multilayer_feature
+        self.upstream, self.featurizer = self._get_upstream(frontend_conf)
+        self.pretrained_params = copy.deepcopy(self.upstream.state_dict())
+        self.output_dim = self.featurizer.output_dim
+        self.frontend_type = "s3prl"
+        self.hop_length = self.upstream.get_downsample_rates("key")
+
+    def _get_upstream(self, frontend_conf):
+        """Get S3PRL upstream model."""
+        s3prl_args = base_s3prl_setup(
+            Namespace(**frontend_conf, device="cpu"),
+        )
+        self.args = s3prl_args
+
+        s3prl_path = None
+        python_path_list = os.environ.get("PYTHONPATH", "(None)").split(":")
+        for p in python_path_list:
+            if p.endswith("s3prl"):
+                s3prl_path = p
+                break
+        assert s3prl_path is not None
+
+        s3prl_upstream = torch.hub.load(
+            s3prl_path,
+            s3prl_args.upstream,
+            ckpt=s3prl_args.upstream_ckpt,
+            model_config=s3prl_args.upstream_model_config,
+            refresh=s3prl_args.upstream_refresh,
+            source="local",
+        ).to("cpu")
+
+        if getattr(
+            s3prl_upstream, "model", None
+        ) is not None and s3prl_upstream.model.__class__.__name__ in [
+            "Wav2Vec2Model",
+            "HubertModel",
+        ]:
+            s3prl_upstream.model.encoder.layerdrop = 0.0
+
+        from s3prl.upstream.interfaces import Featurizer
+
+        if self.multilayer_feature is None:
+            feature_selection = "last_hidden_state"
+        else:
+            feature_selection = "hidden_states"
+        s3prl_featurizer = Featurizer(
+            upstream=s3prl_upstream,
+            feature_selection=feature_selection,
+            upstream_device="cpu",
+        )
+
+        return s3prl_upstream, s3prl_featurizer
+
+    def _tile_representations(self, feature):
+        """Tile up the representations by `tile_factor`.
+        Input - sequence of representations
+                shape: (batch_size, seq_len, feature_dim)
+        Output - sequence of tiled representations
+                 shape: (batch_size, seq_len * factor, feature_dim)
+        """
+        assert len(feature.shape) == 3, "Input argument `feature` has invalid shape: {}".format(
+            feature.shape
+        )
+        tiled_feature = feature.repeat(1, 1, self.args.tile_factor)
+        tiled_feature = tiled_feature.reshape(
+            feature.size(0), feature.size(1) * self.args.tile_factor, feature.size(2)
+        )
+        return tiled_feature
+
+    def output_size(self) -> int:
+        """Output size."""
+        return self.output_dim
+
+    def forward(
+        self, input: torch.Tensor, input_lengths: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass for training.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+            """
+        wavs = [wav[: input_lengths[i]] for i, wav in enumerate(input)]
+        self.upstream.eval()
+        with torch.no_grad():
+            feats = self.upstream(wavs)
+        feats = self.featurizer(wavs, feats)
+
+        if self.args.tile_factor != 1:
+            feats = self._tile_representations(feats)
+
+        input_feats = pad_list(feats, 0.0)
+        feats_lens = torch.tensor([f.shape[0] for f in feats], dtype=torch.long)
+
+        # Saving CUDA Memory
+        del feats
+
+        return input_feats, feats_lens
+
+    def reload_pretrained_parameters(self):
+        """Reload pretrained parameters."""
+        self.upstream.load_state_dict(self.pretrained_params)
+        logging.info("Pretrained S3PRL frontend model parameters reloaded!")

funasr/frontends/utils/__init__.py

@@ -0,0 +1 @@
+"""Initialize sub package."""

funasr/frontends/utils/beamformer.py

@@ -0,0 +1,88 @@
+import torch
+from torch_complex import functional as FC
+from torch_complex.tensor import ComplexTensor
+
+
+def get_power_spectral_density_matrix(
+    xs: ComplexTensor, mask: torch.Tensor, normalization=True, eps: float = 1e-15
+) -> ComplexTensor:
+    """Return cross-channel power spectral density (PSD) matrix
+
+    Args:
+        xs (ComplexTensor): (..., F, C, T)
+        mask (torch.Tensor): (..., F, C, T)
+        normalization (bool):
+        eps (float):
+    Returns
+        psd (ComplexTensor): (..., F, C, C)
+
+    """
+    # outer product: (..., C_1, T) x (..., C_2, T) -> (..., T, C, C_2)
+    psd_Y = FC.einsum("...ct,...et->...tce", [xs, xs.conj()])
+
+    # Averaging mask along C: (..., C, T) -> (..., T)
+    mask = mask.mean(dim=-2)
+
+    # Normalized mask along T: (..., T)
+    if normalization:
+        # If assuming the tensor is padded with zero, the summation along
+        # the time axis is same regardless of the padding length.
+        mask = mask / (mask.sum(dim=-1, keepdim=True) + eps)
+
+    # psd: (..., T, C, C)
+    psd = psd_Y * mask[..., None, None]
+    # (..., T, C, C) -> (..., C, C)
+    psd = psd.sum(dim=-3)
+
+    return psd
+
+
+def get_mvdr_vector(
+    psd_s: ComplexTensor,
+    psd_n: ComplexTensor,
+    reference_vector: torch.Tensor,
+    eps: float = 1e-15,
+) -> ComplexTensor:
+    """Return the MVDR(Minimum Variance Distortionless Response) vector:
+
+        h = (Npsd^-1 @ Spsd) / (Tr(Npsd^-1 @ Spsd)) @ u
+
+    Reference:
+        On optimal frequency-domain multichannel linear filtering
+        for noise reduction; M. Souden et al., 2010;
+        https://ieeexplore.ieee.org/document/5089420
+
+    Args:
+        psd_s (ComplexTensor): (..., F, C, C)
+        psd_n (ComplexTensor): (..., F, C, C)
+        reference_vector (torch.Tensor): (..., C)
+        eps (float):
+    Returns:
+        beamform_vector (ComplexTensor)r: (..., F, C)
+    """
+    # Add eps
+    C = psd_n.size(-1)
+    eye = torch.eye(C, dtype=psd_n.dtype, device=psd_n.device)
+    shape = [1 for _ in range(psd_n.dim() - 2)] + [C, C]
+    eye = eye.view(*shape)
+    psd_n += eps * eye
+
+    # numerator: (..., C_1, C_2) x (..., C_2, C_3) -> (..., C_1, C_3)
+    numerator = FC.einsum("...ec,...cd->...ed", [psd_n.inverse(), psd_s])
+    # ws: (..., C, C) / (...,) -> (..., C, C)
+    ws = numerator / (FC.trace(numerator)[..., None, None] + eps)
+    # h: (..., F, C_1, C_2) x (..., C_2) -> (..., F, C_1)
+    beamform_vector = FC.einsum("...fec,...c->...fe", [ws, reference_vector])
+    return beamform_vector
+
+
+def apply_beamforming_vector(beamform_vector: ComplexTensor, mix: ComplexTensor) -> ComplexTensor:
+    # (..., C) x (..., C, T) -> (..., T)
+    """Apply beamforming vector.
+    
+        Args:
+            beamform_vector: TODO.
+            mix: TODO.
+        """
+    es = FC.einsum("...c,...ct->...t", [beamform_vector.conj(), mix])
+    return es

funasr/frontends/utils/complex_utils.py

@@ -0,0 +1,259 @@
+"""Beamformer module."""
+
+from distutils.version import LooseVersion
+from typing import Sequence
+from typing import Tuple
+from typing import Union
+
+import torch
+
+try:
+    from torch_complex import functional as FC
+    from torch_complex.tensor import ComplexTensor
+except:
+    print("Please install torch_complex firstly")
+
+
+EPS = torch.finfo(torch.double).eps
+is_torch_1_8_plus = LooseVersion(torch.__version__) >= LooseVersion("1.8.0")
+is_torch_1_9_plus = LooseVersion(torch.__version__) >= LooseVersion("1.9.0")
+
+
+def new_complex_like(
+    ref: Union[torch.Tensor, ComplexTensor],
+    real_imag: Tuple[torch.Tensor, torch.Tensor],
+):
+    """New complex like.
+    
+        Args:
+            ref: TODO.
+            real_imag: TODO.
+        """
+    if isinstance(ref, ComplexTensor):
+        return ComplexTensor(*real_imag)
+    elif is_torch_complex_tensor(ref):
+        return torch.complex(*real_imag)
+    else:
+        raise ValueError("Please update your PyTorch version to 1.9+ for complex support.")
+
+
+def is_torch_complex_tensor(c):
+    """Is torch complex tensor.
+    
+        Args:
+            c: TODO.
+        """
+    return not isinstance(c, ComplexTensor) and is_torch_1_9_plus and torch.is_complex(c)
+
+
+def is_complex(c):
+    """Is complex.
+    
+        Args:
+            c: TODO.
+        """
+    return isinstance(c, ComplexTensor) or is_torch_complex_tensor(c)
+
+
+def to_double(c):
+    """To double.
+    
+        Args:
+            c: TODO.
+        """
+    if not isinstance(c, ComplexTensor) and is_torch_1_9_plus and torch.is_complex(c):
+        return c.to(dtype=torch.complex128)
+    else:
+        return c.double()
+
+
+def to_float(c):
+    """To float.
+    
+        Args:
+            c: TODO.
+        """
+    if not isinstance(c, ComplexTensor) and is_torch_1_9_plus and torch.is_complex(c):
+        return c.to(dtype=torch.complex64)
+    else:
+        return c.float()
+
+
+def cat(seq: Sequence[Union[ComplexTensor, torch.Tensor]], *args, **kwargs):
+    """Cat.
+    
+        Args:
+            seq: TODO.
+            *args: Variable positional arguments.
+            **kwargs: Additional keyword arguments.
+        """
+    if not isinstance(seq, (list, tuple)):
+        raise TypeError(
+            "cat(): argument 'tensors' (position 1) must be tuple of Tensors, " "not Tensor"
+        )
+    if isinstance(seq[0], ComplexTensor):
+        return FC.cat(seq, *args, **kwargs)
+    else:
+        return torch.cat(seq, *args, **kwargs)
+
+
+def complex_norm(c: Union[torch.Tensor, ComplexTensor], dim=-1, keepdim=False) -> torch.Tensor:
+    """Complex norm.
+    
+        Args:
+            c: TODO.
+            dim: TODO.
+            keepdim: TODO.
+        """
+    if not is_complex(c):
+        raise TypeError("Input is not a complex tensor.")
+    if is_torch_complex_tensor(c):
+        return torch.norm(c, dim=dim, keepdim=keepdim)
+    else:
+        return torch.sqrt((c.real**2 + c.imag**2).sum(dim=dim, keepdim=keepdim) + EPS)
+
+
+def einsum(equation, *operands):
+    # NOTE: Do not mix ComplexTensor and torch.complex in the input!
+    # NOTE (wangyou): Until PyTorch 1.9.0, torch.einsum does not support
+    # mixed input with complex and real tensors.
+    """Einsum.
+    
+        Args:
+            equation: TODO.
+            *operands: Variable positional arguments.
+        """
+    if len(operands) == 1:
+        if isinstance(operands[0], (tuple, list)):
+            operands = operands[0]
+        complex_module = FC if isinstance(operands[0], ComplexTensor) else torch
+        return complex_module.einsum(equation, *operands)
+    elif len(operands) != 2:
+        op0 = operands[0]
+        same_type = all(op.dtype == op0.dtype for op in operands[1:])
+        if same_type:
+            _einsum = FC.einsum if isinstance(op0, ComplexTensor) else torch.einsum
+            return _einsum(equation, *operands)
+        else:
+            raise ValueError("0 or More than 2 operands are not supported.")
+    a, b = operands
+    if isinstance(a, ComplexTensor) or isinstance(b, ComplexTensor):
+        return FC.einsum(equation, a, b)
+    elif is_torch_1_9_plus and (torch.is_complex(a) or torch.is_complex(b)):
+        if not torch.is_complex(a):
+            o_real = torch.einsum(equation, a, b.real)
+            o_imag = torch.einsum(equation, a, b.imag)
+            return torch.complex(o_real, o_imag)
+        elif not torch.is_complex(b):
+            o_real = torch.einsum(equation, a.real, b)
+            o_imag = torch.einsum(equation, a.imag, b)
+            return torch.complex(o_real, o_imag)
+        else:
+            return torch.einsum(equation, a, b)
+    else:
+        return torch.einsum(equation, a, b)
+
+
+def inverse(c: Union[torch.Tensor, ComplexTensor]) -> Union[torch.Tensor, ComplexTensor]:
+    """Inverse.
+    
+        Args:
+            c: TODO.
+        """
+    if isinstance(c, ComplexTensor):
+        return c.inverse2()
+    else:
+        return c.inverse()
+
+
+def matmul(
+    a: Union[torch.Tensor, ComplexTensor], b: Union[torch.Tensor, ComplexTensor]
+) -> Union[torch.Tensor, ComplexTensor]:
+    # NOTE: Do not mix ComplexTensor and torch.complex in the input!
+    # NOTE (wangyou): Until PyTorch 1.9.0, torch.matmul does not support
+    # multiplication between complex and real tensors.
+    """Matmul.
+    
+        Args:
+            a: TODO.
+            b: TODO.
+        """
+    if isinstance(a, ComplexTensor) or isinstance(b, ComplexTensor):
+        return FC.matmul(a, b)
+    elif is_torch_1_9_plus and (torch.is_complex(a) or torch.is_complex(b)):
+        if not torch.is_complex(a):
+            o_real = torch.matmul(a, b.real)
+            o_imag = torch.matmul(a, b.imag)
+            return torch.complex(o_real, o_imag)
+        elif not torch.is_complex(b):
+            o_real = torch.matmul(a.real, b)
+            o_imag = torch.matmul(a.imag, b)
+            return torch.complex(o_real, o_imag)
+        else:
+            return torch.matmul(a, b)
+    else:
+        return torch.matmul(a, b)
+
+
+def trace(a: Union[torch.Tensor, ComplexTensor]):
+    # NOTE (wangyou): until PyTorch 1.9.0, torch.trace does not
+    # support bacth processing. Use FC.trace() as fallback.
+    """Trace.
+    
+        Args:
+            a: TODO.
+        """
+    return FC.trace(a)
+
+
+def reverse(a: Union[torch.Tensor, ComplexTensor], dim=0):
+    """Reverse.
+    
+        Args:
+            a: TODO.
+            dim: TODO.
+        """
+    if isinstance(a, ComplexTensor):
+        return FC.reverse(a, dim=dim)
+    else:
+        return torch.flip(a, dims=(dim,))
+
+
+def solve(b: Union[torch.Tensor, ComplexTensor], a: Union[torch.Tensor, ComplexTensor]):
+    """Solve the linear equation ax = b."""
+    # NOTE: Do not mix ComplexTensor and torch.complex in the input!
+    # NOTE (wangyou): Until PyTorch 1.9.0, torch.solve does not support
+    # mixed input with complex and real tensors.
+    if isinstance(a, ComplexTensor) or isinstance(b, ComplexTensor):
+        if isinstance(a, ComplexTensor) and isinstance(b, ComplexTensor):
+            return FC.solve(b, a, return_LU=False)
+        else:
+            return matmul(inverse(a), b)
+    elif is_torch_1_9_plus and (torch.is_complex(a) or torch.is_complex(b)):
+        if torch.is_complex(a) and torch.is_complex(b):
+            return torch.linalg.solve(a, b)
+        else:
+            return matmul(inverse(a), b)
+    else:
+        if is_torch_1_8_plus:
+            return torch.linalg.solve(a, b)
+        else:
+            return torch.solve(b, a)[0]
+
+
+def stack(seq: Sequence[Union[ComplexTensor, torch.Tensor]], *args, **kwargs):
+    """Stack.
+    
+        Args:
+            seq: TODO.
+            *args: Variable positional arguments.
+            **kwargs: Additional keyword arguments.
+        """
+    if not isinstance(seq, (list, tuple)):
+        raise TypeError(
+            "stack(): argument 'tensors' (position 1) must be tuple of Tensors, " "not Tensor"
+        )
+    if isinstance(seq[0], ComplexTensor):
+        return FC.stack(seq, *args, **kwargs)
+    else:
+        return torch.stack(seq, *args, **kwargs)

funasr/frontends/utils/dnn_beamformer.py

@@ -0,0 +1,189 @@
+"""DNN beamformer module."""
+
+from typing import Tuple
+
+import torch
+from torch.nn import functional as F
+
+from funasr.frontends.utils.beamformer import apply_beamforming_vector
+from funasr.frontends.utils.beamformer import get_mvdr_vector
+from funasr.frontends.utils.beamformer import (
+    get_power_spectral_density_matrix,  # noqa: H301
+)
+from funasr.frontends.utils.mask_estimator import MaskEstimator
+from torch_complex.tensor import ComplexTensor
+
+
+class DNN_Beamformer(torch.nn.Module):
+    """DNN mask based Beamformer
+
+    Citation:
+        Multichannel End-to-end Speech Recognition; T. Ochiai et al., 2017;
+        https://arxiv.org/abs/1703.04783
+
+    """
+
+    def __init__(
+        self,
+        bidim,
+        btype="blstmp",
+        blayers=3,
+        bunits=300,
+        bprojs=320,
+        bnmask=2,
+        dropout_rate=0.0,
+        badim=320,
+        ref_channel: int = -1,
+        beamformer_type="mvdr",
+    ):
+        """Initialize DNN_Beamformer.
+        
+            Args:
+                bidim: TODO.
+                btype: TODO.
+                blayers: TODO.
+                bunits: TODO.
+                bprojs: TODO.
+                bnmask: TODO.
+                dropout_rate: TODO.
+                badim: TODO.
+                ref_channel: TODO.
+                beamformer_type: TODO.
+            """
+        super().__init__()
+        self.mask = MaskEstimator(btype, bidim, blayers, bunits, bprojs, dropout_rate, nmask=bnmask)
+        self.ref = AttentionReference(bidim, badim)
+        self.ref_channel = ref_channel
+
+        self.nmask = bnmask
+
+        if beamformer_type != "mvdr":
+            raise ValueError("Not supporting beamformer_type={}".format(beamformer_type))
+        self.beamformer_type = beamformer_type
+
+    def forward(
+        self, data: ComplexTensor, ilens: torch.LongTensor
+    ) -> Tuple[ComplexTensor, torch.LongTensor, ComplexTensor]:
+        """The forward function
+
+        Notation:
+            B: Batch
+            C: Channel
+            T: Time or Sequence length
+            F: Freq
+
+        Args:
+            data (ComplexTensor): (B, T, C, F)
+            ilens (torch.Tensor): (B,)
+        Returns:
+            enhanced (ComplexTensor): (B, T, F)
+            ilens (torch.Tensor): (B,)
+
+        """
+
+        def apply_beamforming(data, ilens, psd_speech, psd_noise):
+            # u: (B, C)
+            """Apply beamforming.
+            
+                Args:
+                    data: TODO.
+                    ilens: TODO.
+                    psd_speech: TODO.
+                    psd_noise: TODO.
+                """
+            if self.ref_channel < 0:
+                u, _ = self.ref(psd_speech, ilens)
+            else:
+                # (optional) Create onehot vector for fixed reference microphone
+                u = torch.zeros(*(data.size()[:-3] + (data.size(-2),)), device=data.device)
+                u[..., self.ref_channel].fill_(1)
+
+            ws = get_mvdr_vector(psd_speech, psd_noise, u)
+            enhanced = apply_beamforming_vector(ws, data)
+
+            return enhanced, ws
+
+        # data (B, T, C, F) -> (B, F, C, T)
+        data = data.permute(0, 3, 2, 1)
+
+        # mask: (B, F, C, T)
+        masks, _ = self.mask(data, ilens)
+        assert self.nmask == len(masks)
+
+        if self.nmask == 2:  # (mask_speech, mask_noise)
+            mask_speech, mask_noise = masks
+
+            psd_speech = get_power_spectral_density_matrix(data, mask_speech)
+            psd_noise = get_power_spectral_density_matrix(data, mask_noise)
+
+            enhanced, ws = apply_beamforming(data, ilens, psd_speech, psd_noise)
+
+            # (..., F, T) -> (..., T, F)
+            enhanced = enhanced.transpose(-1, -2)
+            mask_speech = mask_speech.transpose(-1, -3)
+        else:  # multi-speaker case: (mask_speech1, ..., mask_noise)
+            mask_speech = list(masks[:-1])
+            mask_noise = masks[-1]
+
+            psd_speeches = [get_power_spectral_density_matrix(data, mask) for mask in mask_speech]
+            psd_noise = get_power_spectral_density_matrix(data, mask_noise)
+
+            enhanced = []
+            ws = []
+            for i in range(self.nmask - 1):
+                psd_speech = psd_speeches.pop(i)
+                # treat all other speakers' psd_speech as noises
+                enh, w = apply_beamforming(data, ilens, psd_speech, sum(psd_speeches) + psd_noise)
+                psd_speeches.insert(i, psd_speech)
+
+                # (..., F, T) -> (..., T, F)
+                enh = enh.transpose(-1, -2)
+                mask_speech[i] = mask_speech[i].transpose(-1, -3)
+
+                enhanced.append(enh)
+                ws.append(w)
+
+        return enhanced, ilens, mask_speech
+
+
+class AttentionReference(torch.nn.Module):
+    def __init__(self, bidim, att_dim):
+        """Initialize AttentionReference.
+        
+            Args:
+                bidim: TODO.
+                att_dim: Size/dimension parameter.
+            """
+        super().__init__()
+        self.mlp_psd = torch.nn.Linear(bidim, att_dim)
+        self.gvec = torch.nn.Linear(att_dim, 1)
+
+    def forward(
+        self, psd_in: ComplexTensor, ilens: torch.LongTensor, scaling: float = 2.0
+    ) -> Tuple[torch.Tensor, torch.LongTensor]:
+        """The forward function
+
+        Args:
+            psd_in (ComplexTensor): (B, F, C, C)
+            ilens (torch.Tensor): (B,)
+            scaling (float):
+        Returns:
+            u (torch.Tensor): (B, C)
+            ilens (torch.Tensor): (B,)
+        """
+        B, _, C = psd_in.size()[:3]
+        assert psd_in.size(2) == psd_in.size(3), psd_in.size()
+        # psd_in: (B, F, C, C)
+        psd = psd_in.masked_fill(torch.eye(C, dtype=torch.bool, device=psd_in.device), 0)
+        # psd: (B, F, C, C) -> (B, C, F)
+        psd = (psd.sum(dim=-1) / (C - 1)).transpose(-1, -2)
+
+        # Calculate amplitude
+        psd_feat = (psd.real**2 + psd.imag**2) ** 0.5
+
+        # (B, C, F) -> (B, C, F2)
+        mlp_psd = self.mlp_psd(psd_feat)
+        # (B, C, F2) -> (B, C, 1) -> (B, C)
+        e = self.gvec(torch.tanh(mlp_psd)).squeeze(-1)
+        u = F.softmax(scaling * e, dim=-1)
+        return u, ilens

funasr/frontends/utils/dnn_wpe.py

@@ -0,0 +1,108 @@
+from typing import Tuple
+
+from pytorch_wpe import wpe_one_iteration
+import torch
+from torch_complex.tensor import ComplexTensor
+
+from funasr.frontends.utils.mask_estimator import MaskEstimator
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+
+
+class DNN_WPE(torch.nn.Module):
+    def __init__(
+        self,
+        wtype: str = "blstmp",
+        widim: int = 257,
+        wlayers: int = 3,
+        wunits: int = 300,
+        wprojs: int = 320,
+        dropout_rate: float = 0.0,
+        taps: int = 5,
+        delay: int = 3,
+        use_dnn_mask: bool = True,
+        iterations: int = 1,
+        normalization: bool = False,
+    ):
+        """Initialize DNN_WPE.
+        
+            Args:
+                wtype: TODO.
+                widim: TODO.
+                wlayers: TODO.
+                wunits: TODO.
+                wprojs: TODO.
+                dropout_rate: TODO.
+                taps: TODO.
+                delay: TODO.
+                use_dnn_mask: TODO.
+                iterations: TODO.
+                normalization: TODO.
+            """
+        super().__init__()
+        self.iterations = iterations
+        self.taps = taps
+        self.delay = delay
+
+        self.normalization = normalization
+        self.use_dnn_mask = use_dnn_mask
+
+        self.inverse_power = True
+
+        if self.use_dnn_mask:
+            self.mask_est = MaskEstimator(
+                wtype, widim, wlayers, wunits, wprojs, dropout_rate, nmask=1
+            )
+
+    def forward(
+        self, data: ComplexTensor, ilens: torch.LongTensor
+    ) -> Tuple[ComplexTensor, torch.LongTensor, ComplexTensor]:
+        """The forward function
+
+        Notation:
+            B: Batch
+            C: Channel
+            T: Time or Sequence length
+            F: Freq or Some dimension of the feature vector
+
+        Args:
+            data: (B, C, T, F)
+            ilens: (B,)
+        Returns:
+            data: (B, C, T, F)
+            ilens: (B,)
+        """
+        # (B, T, C, F) -> (B, F, C, T)
+        enhanced = data = data.permute(0, 3, 2, 1)
+        mask = None
+
+        for i in range(self.iterations):
+            # Calculate power: (..., C, T)
+            power = enhanced.real**2 + enhanced.imag**2
+            if i == 0 and self.use_dnn_mask:
+                # mask: (B, F, C, T)
+                (mask,), _ = self.mask_est(enhanced, ilens)
+                if self.normalization:
+                    # Normalize along T
+                    mask = mask / mask.sum(dim=-1)[..., None]
+                # (..., C, T) * (..., C, T) -> (..., C, T)
+                power = power * mask
+
+            # Averaging along the channel axis: (..., C, T) -> (..., T)
+            power = power.mean(dim=-2)
+
+            # enhanced: (..., C, T) -> (..., C, T)
+            enhanced = wpe_one_iteration(
+                data.contiguous(),
+                power,
+                taps=self.taps,
+                delay=self.delay,
+                inverse_power=self.inverse_power,
+            )
+
+            enhanced.masked_fill_(make_pad_mask(ilens, enhanced.real), 0)
+
+        # (B, F, C, T) -> (B, T, C, F)
+        enhanced = enhanced.permute(0, 3, 2, 1)
+        if mask is not None:
+            mask = mask.transpose(-1, -3)
+        return enhanced, ilens, mask

funasr/frontends/utils/feature_transform.py

@@ -0,0 +1,331 @@
+from typing import List
+from typing import Tuple
+from typing import Union
+
+import librosa
+import numpy as np
+import torch
+from torch_complex.tensor import ComplexTensor
+
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+
+
+class FeatureTransform(torch.nn.Module):
+    def __init__(
+        self,
+        # Mel options,
+        fs: int = 16000,
+        n_fft: int = 512,
+        n_mels: int = 80,
+        fmin: float = 0.0,
+        fmax: float = None,
+        # Normalization
+        stats_file: str = None,
+        apply_uttmvn: bool = True,
+        uttmvn_norm_means: bool = True,
+        uttmvn_norm_vars: bool = False,
+    ):
+        """Initialize FeatureTransform.
+        
+            Args:
+                fs: TODO.
+                n_fft: TODO.
+                n_mels: TODO.
+                fmin: TODO.
+                fmax: TODO.
+                stats_file: TODO.
+                apply_uttmvn: TODO.
+                uttmvn_norm_means: TODO.
+                uttmvn_norm_vars: TODO.
+            """
+        super().__init__()
+        self.apply_uttmvn = apply_uttmvn
+
+        self.logmel = LogMel(fs=fs, n_fft=n_fft, n_mels=n_mels, fmin=fmin, fmax=fmax)
+        self.stats_file = stats_file
+        if stats_file is not None:
+            self.global_mvn = GlobalMVN(stats_file)
+        else:
+            self.global_mvn = None
+
+        if self.apply_uttmvn is not None:
+            self.uttmvn = UtteranceMVN(norm_means=uttmvn_norm_means, norm_vars=uttmvn_norm_vars)
+        else:
+            self.uttmvn = None
+
+    def forward(
+        self, x: ComplexTensor, ilens: Union[torch.LongTensor, np.ndarray, List[int]]
+    ) -> Tuple[torch.Tensor, torch.LongTensor]:
+        # (B, T, F) or (B, T, C, F)
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+                ilens: TODO.
+            """
+        if x.dim() not in (3, 4):
+            raise ValueError(f"Input dim must be 3 or 4: {x.dim()}")
+        if not torch.is_tensor(ilens):
+            ilens = torch.from_numpy(np.asarray(ilens)).to(x.device)
+
+        if x.dim() == 4:
+            # h: (B, T, C, F) -> h: (B, T, F)
+            if self.training:
+                # Select 1ch randomly
+                ch = np.random.randint(x.size(2))
+                h = x[:, :, ch, :]
+            else:
+                # Use the first channel
+                h = x[:, :, 0, :]
+        else:
+            h = x
+
+        # h: ComplexTensor(B, T, F) -> torch.Tensor(B, T, F)
+        h = h.real**2 + h.imag**2
+
+        h, _ = self.logmel(h, ilens)
+        if self.stats_file is not None:
+            h, _ = self.global_mvn(h, ilens)
+        if self.apply_uttmvn:
+            h, _ = self.uttmvn(h, ilens)
+
+        return h, ilens
+
+
+class LogMel(torch.nn.Module):
+    """Convert STFT to fbank feats
+
+    The arguments is same as librosa.filters.mel
+
+    Args:
+        fs: number > 0 [scalar] sampling rate of the incoming signal
+        n_fft: int > 0 [scalar] number of FFT components
+        n_mels: int > 0 [scalar] number of Mel bands to generate
+        fmin: float >= 0 [scalar] lowest frequency (in Hz)
+        fmax: float >= 0 [scalar] highest frequency (in Hz).
+            If `None`, use `fmax = fs / 2.0`
+        htk: use HTK formula instead of Slaney
+        norm: {None, 1, np.inf} [scalar]
+            if 1, divide the triangular mel weights by the width of the mel band
+            (area normalization).  Otherwise, leave all the triangles aiming for
+            a peak value of 1.0
+
+    """
+
+    def __init__(
+        self,
+        fs: int = 16000,
+        n_fft: int = 512,
+        n_mels: int = 80,
+        fmin: float = 0.0,
+        fmax: float = None,
+        htk: bool = False,
+        norm=1,
+    ):
+        """Initialize LogMel.
+        
+            Args:
+                fs: TODO.
+                n_fft: TODO.
+                n_mels: TODO.
+                fmin: TODO.
+                fmax: TODO.
+                htk: TODO.
+                norm: TODO.
+            """
+        super().__init__()
+
+        _mel_options = dict(
+            sr=fs, n_fft=n_fft, n_mels=n_mels, fmin=fmin, fmax=fmax, htk=htk, norm=norm
+        )
+        self.mel_options = _mel_options
+
+        # Note(kamo): The mel matrix of librosa is different from kaldi.
+        melmat = librosa.filters.mel(**_mel_options)
+        # melmat: (D2, D1) -> (D1, D2)
+        self.register_buffer("melmat", torch.from_numpy(melmat.T).float())
+
+    def extra_repr(self):
+        """Extra repr."""
+        return ", ".join(f"{k}={v}" for k, v in self.mel_options.items())
+
+    def forward(
+        self, feat: torch.Tensor, ilens: torch.LongTensor
+    ) -> Tuple[torch.Tensor, torch.LongTensor]:
+        # feat: (B, T, D1) x melmat: (D1, D2) -> mel_feat: (B, T, D2)
+        """Forward pass for training.
+        
+            Args:
+                feat: TODO.
+                ilens: TODO.
+            """
+        mel_feat = torch.matmul(feat, self.melmat)
+
+        logmel_feat = (mel_feat + 1e-20).log()
+        # Zero padding
+        logmel_feat = logmel_feat.masked_fill(make_pad_mask(ilens, logmel_feat, 1), 0.0)
+        return logmel_feat, ilens
+
+
+class GlobalMVN(torch.nn.Module):
+    """Apply global mean and variance normalization
+
+    Args:
+        stats_file(str): npy file of 1-dim array or text file.
+            From the _first element to
+            the {(len(array) - 1) / 2}th element are treated as
+            the sum of features,
+            and the rest excluding the last elements are
+            treated as the sum of the square value of features,
+            and the last elements eqauls to the number of samples.
+        std_floor(float):
+    """
+
+    def __init__(
+        self,
+        stats_file: str,
+        norm_means: bool = True,
+        norm_vars: bool = True,
+        eps: float = 1.0e-20,
+    ):
+        """Initialize GlobalMVN.
+        
+            Args:
+                stats_file: TODO.
+                norm_means: TODO.
+                norm_vars: TODO.
+                eps: TODO.
+            """
+        super().__init__()
+        self.norm_means = norm_means
+        self.norm_vars = norm_vars
+
+        self.stats_file = stats_file
+        stats = np.load(stats_file)
+
+        stats = stats.astype(float)
+        assert (len(stats) - 1) % 2 == 0, stats.shape
+
+        count = stats.flatten()[-1]
+        mean = stats[: (len(stats) - 1) // 2] / count
+        var = stats[(len(stats) - 1) // 2 : -1] / count - mean * mean
+        std = np.maximum(np.sqrt(var), eps)
+
+        self.register_buffer("bias", torch.from_numpy(-mean.astype(np.float32)))
+        self.register_buffer("scale", torch.from_numpy(1 / std.astype(np.float32)))
+
+    def extra_repr(self):
+        """Extra repr."""
+        return (
+            f"stats_file={self.stats_file}, "
+            f"norm_means={self.norm_means}, norm_vars={self.norm_vars}"
+        )
+
+    def forward(
+        self, x: torch.Tensor, ilens: torch.LongTensor
+    ) -> Tuple[torch.Tensor, torch.LongTensor]:
+        # feat: (B, T, D)
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+                ilens: TODO.
+            """
+        if self.norm_means:
+            x += self.bias.type_as(x)
+            x.masked_fill(make_pad_mask(ilens, x, 1), 0.0)
+
+        if self.norm_vars:
+            x *= self.scale.type_as(x)
+        return x, ilens
+
+
+class UtteranceMVN(torch.nn.Module):
+    def __init__(self, norm_means: bool = True, norm_vars: bool = False, eps: float = 1.0e-20):
+        """Initialize UtteranceMVN.
+        
+            Args:
+                norm_means: TODO.
+                norm_vars: TODO.
+                eps: TODO.
+            """
+        super().__init__()
+        self.norm_means = norm_means
+        self.norm_vars = norm_vars
+        self.eps = eps
+
+    def extra_repr(self):
+        """Extra repr."""
+        return f"norm_means={self.norm_means}, norm_vars={self.norm_vars}"
+
+    def forward(
+        self, x: torch.Tensor, ilens: torch.LongTensor
+    ) -> Tuple[torch.Tensor, torch.LongTensor]:
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+                ilens: TODO.
+            """
+        return utterance_mvn(
+            x, ilens, norm_means=self.norm_means, norm_vars=self.norm_vars, eps=self.eps
+        )
+
+
+def utterance_mvn(
+    x: torch.Tensor,
+    ilens: torch.LongTensor,
+    norm_means: bool = True,
+    norm_vars: bool = False,
+    eps: float = 1.0e-20,
+) -> Tuple[torch.Tensor, torch.LongTensor]:
+    """Apply utterance mean and variance normalization
+
+    Args:
+        x: (B, T, D), assumed zero padded
+        ilens: (B, T, D)
+        norm_means:
+        norm_vars:
+        eps:
+
+    """
+    ilens_ = ilens.type_as(x)
+    # mean: (B, D)
+    mean = x.sum(dim=1) / ilens_[:, None]
+
+    if norm_means:
+        x -= mean[:, None, :]
+        x_ = x
+    else:
+        x_ = x - mean[:, None, :]
+
+    # Zero padding
+    x_.masked_fill(make_pad_mask(ilens, x_, 1), 0.0)
+    if norm_vars:
+        var = x_.pow(2).sum(dim=1) / ilens_[:, None]
+        var = torch.clamp(var, min=eps)
+        x /= var.sqrt()[:, None, :]
+        x_ = x
+    return x_, ilens
+
+
+def feature_transform_for(args, n_fft):
+    """Feature transform for.
+    
+        Args:
+            args: TODO.
+            n_fft: TODO.
+        """
+    return FeatureTransform(
+        # Mel options,
+        fs=args.fbank_fs,
+        n_fft=n_fft,
+        n_mels=args.n_mels,
+        fmin=args.fbank_fmin,
+        fmax=args.fbank_fmax,
+        # Normalization
+        stats_file=args.stats_file,
+        apply_uttmvn=args.apply_uttmvn,
+        uttmvn_norm_means=args.uttmvn_norm_means,
+        uttmvn_norm_vars=args.uttmvn_norm_vars,
+    )

funasr/frontends/utils/frontend.py

@@ -0,0 +1,186 @@
+from typing import List
+from typing import Optional
+from typing import Tuple
+from typing import Union
+
+import numpy
+import torch
+import torch.nn as nn
+from torch_complex.tensor import ComplexTensor
+
+from funasr.frontends.utils.dnn_beamformer import DNN_Beamformer
+from funasr.frontends.utils.dnn_wpe import DNN_WPE
+
+
+class Frontend(nn.Module):
+    def __init__(
+        self,
+        idim: int,
+        # WPE options
+        use_wpe: bool = False,
+        wtype: str = "blstmp",
+        wlayers: int = 3,
+        wunits: int = 300,
+        wprojs: int = 320,
+        wdropout_rate: float = 0.0,
+        taps: int = 5,
+        delay: int = 3,
+        use_dnn_mask_for_wpe: bool = True,
+        # Beamformer options
+        use_beamformer: bool = False,
+        btype: str = "blstmp",
+        blayers: int = 3,
+        bunits: int = 300,
+        bprojs: int = 320,
+        bnmask: int = 2,
+        badim: int = 320,
+        ref_channel: int = -1,
+        bdropout_rate=0.0,
+    ):
+        """Initialize Frontend.
+        
+            Args:
+                idim: TODO.
+                use_wpe: TODO.
+                wtype: TODO.
+                wlayers: TODO.
+                wunits: TODO.
+                wprojs: TODO.
+                wdropout_rate: TODO.
+                taps: TODO.
+                delay: TODO.
+                use_dnn_mask_for_wpe: TODO.
+                use_beamformer: TODO.
+                btype: TODO.
+                blayers: TODO.
+                bunits: TODO.
+                bprojs: TODO.
+                bnmask: TODO.
+                badim: TODO.
+                ref_channel: TODO.
+                bdropout_rate: TODO.
+            """
+        super().__init__()
+
+        self.use_beamformer = use_beamformer
+        self.use_wpe = use_wpe
+        self.use_dnn_mask_for_wpe = use_dnn_mask_for_wpe
+        # use frontend for all the data,
+        # e.g. in the case of multi-speaker speech separation
+        self.use_frontend_for_all = bnmask > 2
+
+        if self.use_wpe:
+            if self.use_dnn_mask_for_wpe:
+                # Use DNN for power estimation
+                # (Not observed significant gains)
+                iterations = 1
+            else:
+                # Performing as conventional WPE, without DNN Estimator
+                iterations = 2
+
+            self.wpe = DNN_WPE(
+                wtype=wtype,
+                widim=idim,
+                wunits=wunits,
+                wprojs=wprojs,
+                wlayers=wlayers,
+                taps=taps,
+                delay=delay,
+                dropout_rate=wdropout_rate,
+                iterations=iterations,
+                use_dnn_mask=use_dnn_mask_for_wpe,
+            )
+        else:
+            self.wpe = None
+
+        if self.use_beamformer:
+            self.beamformer = DNN_Beamformer(
+                btype=btype,
+                bidim=idim,
+                bunits=bunits,
+                bprojs=bprojs,
+                blayers=blayers,
+                bnmask=bnmask,
+                dropout_rate=bdropout_rate,
+                badim=badim,
+                ref_channel=ref_channel,
+            )
+        else:
+            self.beamformer = None
+
+    def forward(
+        self, x: ComplexTensor, ilens: Union[torch.LongTensor, numpy.ndarray, List[int]]
+    ) -> Tuple[ComplexTensor, torch.LongTensor, Optional[ComplexTensor]]:
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+                ilens: TODO.
+            """
+        assert len(x) == len(ilens), (len(x), len(ilens))
+        # (B, T, F) or (B, T, C, F)
+        if x.dim() not in (3, 4):
+            raise ValueError(f"Input dim must be 3 or 4: {x.dim()}")
+        if not torch.is_tensor(ilens):
+            ilens = torch.from_numpy(numpy.asarray(ilens)).to(x.device)
+
+        mask = None
+        h = x
+        if h.dim() == 4:
+            if self.training:
+                choices = [(False, False)] if not self.use_frontend_for_all else []
+                if self.use_wpe:
+                    choices.append((True, False))
+
+                if self.use_beamformer:
+                    choices.append((False, True))
+
+                use_wpe, use_beamformer = choices[numpy.random.randint(len(choices))]
+
+            else:
+                use_wpe = self.use_wpe
+                use_beamformer = self.use_beamformer
+
+            # 1. WPE
+            if use_wpe:
+                # h: (B, T, C, F) -> h: (B, T, C, F)
+                h, ilens, mask = self.wpe(h, ilens)
+
+            # 2. Beamformer
+            if use_beamformer:
+                # h: (B, T, C, F) -> h: (B, T, F)
+                h, ilens, mask = self.beamformer(h, ilens)
+
+        return h, ilens, mask
+
+
+def frontend_for(args, idim):
+    """Frontend for.
+    
+        Args:
+            args: TODO.
+            idim: TODO.
+        """
+    return Frontend(
+        idim=idim,
+        # WPE options
+        use_wpe=args.use_wpe,
+        wtype=args.wtype,
+        wlayers=args.wlayers,
+        wunits=args.wunits,
+        wprojs=args.wprojs,
+        wdropout_rate=args.wdropout_rate,
+        taps=args.wpe_taps,
+        delay=args.wpe_delay,
+        use_dnn_mask_for_wpe=args.use_dnn_mask_for_wpe,
+        # Beamformer options
+        use_beamformer=args.use_beamformer,
+        btype=args.btype,
+        blayers=args.blayers,
+        bunits=args.bunits,
+        bprojs=args.bprojs,
+        bnmask=args.bnmask,
+        badim=args.badim,
+        ref_channel=args.ref_channel,
+        bdropout_rate=args.bdropout_rate,
+    )

funasr/frontends/utils/log_mel.py

@@ -0,0 +1,97 @@
+import librosa
+import torch
+from typing import Tuple
+
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+
+
+class LogMel(torch.nn.Module):
+    """Convert STFT to fbank feats
+
+    The arguments is same as librosa.filters.mel
+
+    Args:
+        fs: number > 0 [scalar] sampling rate of the incoming signal
+        n_fft: int > 0 [scalar] number of FFT components
+        n_mels: int > 0 [scalar] number of Mel bands to generate
+        fmin: float >= 0 [scalar] lowest frequency (in Hz)
+        fmax: float >= 0 [scalar] highest frequency (in Hz).
+            If `None`, use `fmax = fs / 2.0`
+        htk: use HTK formula instead of Slaney
+    """
+
+    def __init__(
+        self,
+        fs: int = 16000,
+        n_fft: int = 512,
+        n_mels: int = 80,
+        fmin: float = None,
+        fmax: float = None,
+        htk: bool = False,
+        log_base: float = None,
+    ):
+        """Initialize LogMel.
+        
+            Args:
+                fs: TODO.
+                n_fft: TODO.
+                n_mels: TODO.
+                fmin: TODO.
+                fmax: TODO.
+                htk: TODO.
+                log_base: TODO.
+            """
+        super().__init__()
+
+        fmin = 0 if fmin is None else fmin
+        fmax = fs / 2 if fmax is None else fmax
+        _mel_options = dict(
+            sr=fs,
+            n_fft=n_fft,
+            n_mels=n_mels,
+            fmin=fmin,
+            fmax=fmax,
+            htk=htk,
+        )
+        self.mel_options = _mel_options
+        self.log_base = log_base
+
+        # Note(kamo): The mel matrix of librosa is different from kaldi.
+        melmat = librosa.filters.mel(**_mel_options)
+        # melmat: (D2, D1) -> (D1, D2)
+        self.register_buffer("melmat", torch.from_numpy(melmat.T).float())
+
+    def extra_repr(self):
+        """Extra repr."""
+        return ", ".join(f"{k}={v}" for k, v in self.mel_options.items())
+
+    def forward(
+        self,
+        feat: torch.Tensor,
+        ilens: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # feat: (B, T, D1) x melmat: (D1, D2) -> mel_feat: (B, T, D2)
+        """Forward pass for training.
+        
+            Args:
+                feat: TODO.
+                ilens: TODO.
+            """
+        mel_feat = torch.matmul(feat, self.melmat)
+        mel_feat = torch.clamp(mel_feat, min=1e-10)
+
+        if self.log_base is None:
+            logmel_feat = mel_feat.log()
+        elif self.log_base == 2.0:
+            logmel_feat = mel_feat.log2()
+        elif self.log_base == 10.0:
+            logmel_feat = mel_feat.log10()
+        else:
+            logmel_feat = mel_feat.log() / torch.log(self.log_base)
+
+        # Zero padding
+        if ilens is not None:
+            logmel_feat = logmel_feat.masked_fill(make_pad_mask(ilens, logmel_feat, 1), 0.0)
+        else:
+            ilens = feat.new_full([feat.size(0)], fill_value=feat.size(1), dtype=torch.long)
+        return logmel_feat, ilens

funasr/frontends/utils/mask_estimator.py

@@ -0,0 +1,86 @@
+from typing import Tuple
+
+import numpy as np
+import torch
+from torch.nn import functional as F
+from torch_complex.tensor import ComplexTensor
+
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+from funasr.models.language_model.rnn.encoders import RNN
+from funasr.models.language_model.rnn.encoders import RNNP
+
+
+class MaskEstimator(torch.nn.Module):
+    def __init__(self, type, idim, layers, units, projs, dropout, nmask=1):
+        """Initialize MaskEstimator.
+        
+            Args:
+                type: TODO.
+                idim: TODO.
+                layers: TODO.
+                units: TODO.
+                projs: TODO.
+                dropout: TODO.
+                nmask: TODO.
+            """
+        super().__init__()
+        subsample = np.ones(layers + 1, dtype=np.int32)
+
+        typ = type.lstrip("vgg").rstrip("p")
+        if type[-1] == "p":
+            self.brnn = RNNP(idim, layers, units, projs, subsample, dropout, typ=typ)
+        else:
+            self.brnn = RNN(idim, layers, units, projs, dropout, typ=typ)
+
+        self.type = type
+        self.nmask = nmask
+        self.linears = torch.nn.ModuleList([torch.nn.Linear(projs, idim) for _ in range(nmask)])
+
+    def forward(
+        self, xs: ComplexTensor, ilens: torch.LongTensor
+    ) -> Tuple[Tuple[torch.Tensor, ...], torch.LongTensor]:
+        """The forward function
+
+        Args:
+            xs: (B, F, C, T)
+            ilens: (B,)
+        Returns:
+            hs (torch.Tensor): The hidden vector (B, F, C, T)
+            masks: A tuple of the masks. (B, F, C, T)
+            ilens: (B,)
+        """
+        assert xs.size(0) == ilens.size(0), (xs.size(0), ilens.size(0))
+        _, _, C, input_length = xs.size()
+        # (B, F, C, T) -> (B, C, T, F)
+        xs = xs.permute(0, 2, 3, 1)
+
+        # Calculate amplitude: (B, C, T, F) -> (B, C, T, F)
+        xs = (xs.real**2 + xs.imag**2) ** 0.5
+        # xs: (B, C, T, F) -> xs: (B * C, T, F)
+        xs = xs.contiguous().view(-1, xs.size(-2), xs.size(-1))
+        # ilens: (B,) -> ilens_: (B * C)
+        ilens_ = ilens[:, None].expand(-1, C).contiguous().view(-1)
+
+        # xs: (B * C, T, F) -> xs: (B * C, T, D)
+        xs, _, _ = self.brnn(xs, ilens_)
+        # xs: (B * C, T, D) -> xs: (B, C, T, D)
+        xs = xs.view(-1, C, xs.size(-2), xs.size(-1))
+
+        masks = []
+        for linear in self.linears:
+            # xs: (B, C, T, D) -> mask:(B, C, T, F)
+            mask = linear(xs)
+
+            mask = torch.sigmoid(mask)
+            # Zero padding
+            mask.masked_fill(make_pad_mask(ilens, mask, length_dim=2), 0)
+
+            # (B, C, T, F) -> (B, F, C, T)
+            mask = mask.permute(0, 3, 1, 2)
+
+            # Take cares of multi gpu cases: If input_length > max(ilens)
+            if mask.size(-1) < input_length:
+                mask = F.pad(mask, [0, input_length - mask.size(-1)], value=0)
+            masks.append(mask)
+
+        return tuple(masks), ilens

funasr/frontends/utils/stft.py

@@ -0,0 +1,238 @@
+from distutils.version import LooseVersion
+from typing import Optional
+from typing import Tuple
+from typing import Union
+
+import torch
+
+try:
+    from torch_complex.tensor import ComplexTensor
+except:
+    print("Please install torch_complex firstly")
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+from funasr.frontends.utils.complex_utils import is_complex
+
+import librosa
+import numpy as np
+
+is_torch_1_9_plus = LooseVersion(torch.__version__) >= LooseVersion("1.9.0")
+
+
+is_torch_1_7_plus = LooseVersion(torch.__version__) >= LooseVersion("1.7")
+
+
+class Stft(torch.nn.Module):
+    def __init__(
+        self,
+        n_fft: int = 512,
+        win_length: int = None,
+        hop_length: int = 128,
+        window: Optional[str] = "hann",
+        center: bool = True,
+        normalized: bool = False,
+        onesided: bool = True,
+    ):
+        """Initialize Stft.
+        
+            Args:
+                n_fft: TODO.
+                win_length: TODO.
+                hop_length: TODO.
+                window: TODO.
+                center: TODO.
+                normalized: TODO.
+                onesided: TODO.
+            """
+        super().__init__()
+        self.n_fft = n_fft
+        if win_length is None:
+            self.win_length = n_fft
+        else:
+            self.win_length = win_length
+        self.hop_length = hop_length
+        self.center = center
+        self.normalized = normalized
+        self.onesided = onesided
+        if window is not None and not hasattr(torch, f"{window}_window"):
+            if window.lower() != "povey":
+                raise ValueError(f"{window} window is not implemented")
+        self.window = window
+
+    def extra_repr(self):
+        """Extra repr."""
+        return (
+            f"n_fft={self.n_fft}, "
+            f"win_length={self.win_length}, "
+            f"hop_length={self.hop_length}, "
+            f"center={self.center}, "
+            f"normalized={self.normalized}, "
+            f"onesided={self.onesided}"
+        )
+
+    def forward(
+        self, input: torch.Tensor, ilens: torch.Tensor = None
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """STFT forward function.
+
+        Args:
+            input: (Batch, Nsamples) or (Batch, Nsample, Channels)
+            ilens: (Batch)
+        Returns:
+            output: (Batch, Frames, Freq, 2) or (Batch, Frames, Channels, Freq, 2)
+
+        """
+        bs = input.size(0)
+        if input.dim() == 3:
+            multi_channel = True
+            # input: (Batch, Nsample, Channels) -> (Batch * Channels, Nsample)
+            input = input.transpose(1, 2).reshape(-1, input.size(1))
+        else:
+            multi_channel = False
+
+        # NOTE(kamo):
+        #   The default behaviour of torch.stft is compatible with librosa.stft
+        #   about padding and scaling.
+        #   Note that it's different from scipy.signal.stft
+
+        # output: (Batch, Freq, Frames, 2=real_imag)
+        # or (Batch, Channel, Freq, Frames, 2=real_imag)
+        if self.window is not None:
+            if self.window.lower() == "povey":
+                window = torch.hann_window(
+                    self.win_length, periodic=False, device=input.device, dtype=input.dtype
+                ).pow(0.85)
+            else:
+                window_func = getattr(torch, f"{self.window}_window")
+                window = window_func(self.win_length, dtype=input.dtype, device=input.device)
+        else:
+            window = None
+
+        # For the compatibility of ARM devices, which do not support
+        # torch.stft() due to the lake of MKL.
+        if input.is_cuda or torch.backends.mkl.is_available():
+            stft_kwargs = dict(
+                n_fft=self.n_fft,
+                win_length=self.win_length,
+                hop_length=self.hop_length,
+                center=self.center,
+                window=window,
+                normalized=self.normalized,
+                onesided=self.onesided,
+            )
+            if is_torch_1_7_plus:
+                stft_kwargs["return_complex"] = False
+            output = torch.stft(input, **stft_kwargs)
+        else:
+            if self.training:
+                raise NotImplementedError(
+                    "stft is implemented with librosa on this device, which does not "
+                    "support the training mode."
+                )
+
+            # use stft_kwargs to flexibly control different PyTorch versions' kwargs
+            stft_kwargs = dict(
+                n_fft=self.n_fft,
+                win_length=self.win_length,
+                hop_length=self.hop_length,
+                center=self.center,
+                window=window,
+            )
+
+            if window is not None:
+                # pad the given window to n_fft
+                n_pad_left = (self.n_fft - window.shape[0]) // 2
+                n_pad_right = self.n_fft - window.shape[0] - n_pad_left
+                stft_kwargs["window"] = torch.cat(
+                    [torch.zeros(n_pad_left), window, torch.zeros(n_pad_right)], 0
+                ).numpy()
+            else:
+                win_length = self.win_length if self.win_length is not None else self.n_fft
+                stft_kwargs["window"] = torch.ones(win_length)
+
+            output = []
+            # iterate over istances in a batch
+            for i, instance in enumerate(input):
+                stft = librosa.stft(input[i].numpy(), **stft_kwargs)
+                output.append(torch.tensor(np.stack([stft.real, stft.imag], -1)))
+            output = torch.stack(output, 0)
+            if not self.onesided:
+                len_conj = self.n_fft - output.shape[1]
+                conj = output[:, 1 : 1 + len_conj].flip(1)
+                conj[:, :, :, -1].data *= -1
+                output = torch.cat([output, conj], 1)
+            if self.normalized:
+                output = output * (stft_kwargs["window"].shape[0] ** (-0.5))
+
+        # output: (Batch, Freq, Frames, 2=real_imag)
+        # -> (Batch, Frames, Freq, 2=real_imag)
+        output = output.transpose(1, 2)
+        if multi_channel:
+            # output: (Batch * Channel, Frames, Freq, 2=real_imag)
+            # -> (Batch, Frame, Channel, Freq, 2=real_imag)
+            output = output.view(bs, -1, output.size(1), output.size(2), 2).transpose(1, 2)
+
+        if ilens is not None:
+            if self.center:
+                pad = self.n_fft // 2
+                ilens = ilens + 2 * pad
+
+            olens = (ilens - self.n_fft) // self.hop_length + 1
+            output.masked_fill_(make_pad_mask(olens, output, 1), 0.0)
+        else:
+            olens = None
+
+        return output, olens
+
+    def inverse(
+        self, input: Union[torch.Tensor, ComplexTensor], ilens: torch.Tensor = None
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Inverse STFT.
+
+        Args:
+            input: Tensor(batch, T, F, 2) or ComplexTensor(batch, T, F)
+            ilens: (batch,)
+        Returns:
+            wavs: (batch, samples)
+            ilens: (batch,)
+        """
+        if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
+            istft = torch.functional.istft
+        else:
+            try:
+                import torchaudio
+            except ImportError:
+                raise ImportError("Please install torchaudio>=0.3.0 or use torch>=1.6.0")
+
+            if not hasattr(torchaudio.functional, "istft"):
+                raise ImportError("Please install torchaudio>=0.3.0 or use torch>=1.6.0")
+            istft = torchaudio.functional.istft
+
+        if self.window is not None:
+            window_func = getattr(torch, f"{self.window}_window")
+            if is_complex(input):
+                datatype = input.real.dtype
+            else:
+                datatype = input.dtype
+            window = window_func(self.win_length, dtype=datatype, device=input.device)
+        else:
+            window = None
+
+        if is_complex(input):
+            input = torch.stack([input.real, input.imag], dim=-1)
+        elif input.shape[-1] != 2:
+            raise TypeError("Invalid input type")
+        input = input.transpose(1, 2)
+
+        wavs = istft(
+            input,
+            n_fft=self.n_fft,
+            hop_length=self.hop_length,
+            win_length=self.win_length,
+            window=window,
+            center=self.center,
+            normalized=self.normalized,
+            onesided=self.onesided,
+            length=ilens.max() if ilens is not None else ilens,
+        )
+
+        return wavs, ilens

funasr/frontends/wav_frontend.py

@@ -0,0 +1,672 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+# Part of the implementation is borrowed from espnet/espnet.
+from typing import Tuple
+import copy
+import numpy as np
+import torch
+import torch.nn as nn
+import torchaudio.compliance.kaldi as kaldi
+from torch.nn.utils.rnn import pad_sequence
+
+import funasr.frontends.eend_ola_feature as eend_ola_feature
+from funasr.register import tables
+
+
+def load_cmvn(cmvn_file):
+    """Load cmvn.
+    
+        Args:
+            cmvn_file: TODO.
+        """
+    with open(cmvn_file, "r", encoding="utf-8") as f:
+        lines = f.readlines()
+    means_list = []
+    vars_list = []
+    for i in range(len(lines)):
+        line_item = lines[i].split()
+        if line_item[0] == "<AddShift>":
+            line_item = lines[i + 1].split()
+            if line_item[0] == "<LearnRateCoef>":
+                add_shift_line = line_item[3 : (len(line_item) - 1)]
+                means_list = list(add_shift_line)
+                continue
+        elif line_item[0] == "<Rescale>":
+            line_item = lines[i + 1].split()
+            if line_item[0] == "<LearnRateCoef>":
+                rescale_line = line_item[3 : (len(line_item) - 1)]
+                vars_list = list(rescale_line)
+                continue
+    means = np.array(means_list).astype(np.float32)
+    vars = np.array(vars_list).astype(np.float32)
+    cmvn = np.array([means, vars])
+    cmvn = torch.as_tensor(cmvn, dtype=torch.float32)
+    return cmvn
+
+
+def apply_cmvn(inputs, cmvn):  # noqa
+    """
+    Apply CMVN with mvn data
+    """
+
+    device = inputs.device
+    dtype = inputs.dtype
+    frame, dim = inputs.shape
+
+    means = cmvn[0:1, :dim]
+    vars = cmvn[1:2, :dim]
+    inputs += means.to(device)
+    inputs *= vars.to(device)
+
+    return inputs.type(torch.float32)
+
+
+def apply_lfr(inputs, lfr_m, lfr_n):
+    """Apply lfr.
+    
+        Args:
+            inputs: TODO.
+            lfr_m: TODO.
+            lfr_n: TODO.
+        """
+    LFR_inputs = []
+    T = inputs.shape[0]
+    T_lfr = int(np.ceil(T / lfr_n))
+    left_padding = inputs[0].repeat((lfr_m - 1) // 2, 1)
+    inputs = torch.vstack((left_padding, inputs))
+    T = T + (lfr_m - 1) // 2
+    feat_dim = inputs.shape[-1]
+    strides = (lfr_n * feat_dim, 1)
+    sizes = (T_lfr, lfr_m * feat_dim)
+    last_idx = (T - lfr_m) // lfr_n + 1
+    num_padding = lfr_m - (T - last_idx * lfr_n)
+    if num_padding > 0:
+        num_padding = (2 * lfr_m - 2 * T + (T_lfr - 1 + last_idx) * lfr_n) / 2 * (T_lfr - last_idx)
+        inputs = torch.vstack([inputs] + [inputs[-1:]] * int(num_padding))
+    LFR_outputs = inputs.as_strided(sizes, strides)
+    return LFR_outputs.clone().type(torch.float32)
+
+
+@tables.register("frontend_classes", "wav_frontend")
+@tables.register("frontend_classes", "WavFrontend")
+class WavFrontend(nn.Module):
+    """Conventional frontend structure for ASR."""
+
+    def __init__(
+        self,
+        cmvn_file: str = None,
+        fs: int = 16000,
+        window: str = "hamming",
+        n_mels: int = 80,
+        frame_length: int = 25,
+        frame_shift: int = 10,
+        filter_length_min: int = -1,
+        filter_length_max: int = -1,
+        lfr_m: int = 1,
+        lfr_n: int = 1,
+        dither: float = 1.0,
+        snip_edges: bool = True,
+        upsacle_samples: bool = True,
+        **kwargs,
+    ):
+        """Initialize WavFrontend.
+        
+            Args:
+                cmvn_file: TODO.
+                fs: TODO.
+                window: TODO.
+                n_mels: TODO.
+                frame_length: TODO.
+                frame_shift: TODO.
+                filter_length_min: TODO.
+                filter_length_max: TODO.
+                lfr_m: TODO.
+                lfr_n: TODO.
+                dither: TODO.
+                snip_edges: TODO.
+                upsacle_samples: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        self.fs = fs
+        self.window = window
+        self.n_mels = n_mels
+        self.frame_length = frame_length
+        self.frame_shift = frame_shift
+        self.filter_length_min = filter_length_min
+        self.filter_length_max = filter_length_max
+        self.lfr_m = lfr_m
+        self.lfr_n = lfr_n
+        self.cmvn_file = cmvn_file
+        self.dither = dither
+        self.snip_edges = snip_edges
+        self.upsacle_samples = upsacle_samples
+        self.cmvn = None if self.cmvn_file is None else load_cmvn(self.cmvn_file)
+
+    def output_size(self) -> int:
+        """Output size."""
+        return self.n_mels * self.lfr_m
+
+    def forward(
+        self,
+        input: torch.Tensor,
+        input_lengths,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass for training.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+                **kwargs: Additional keyword arguments.
+            """
+        batch_size = input.size(0)
+        feats = []
+        feats_lens = []
+        for i in range(batch_size):
+            waveform_length = input_lengths[i]
+            waveform = input[i][:waveform_length]
+            if self.upsacle_samples:
+                waveform = waveform * (1 << 15)
+            waveform = waveform.unsqueeze(0)
+            mat = kaldi.fbank(
+                waveform,
+                num_mel_bins=self.n_mels,
+                frame_length=min(self.frame_length,waveform_length/self.fs*1000),
+                frame_shift=self.frame_shift,
+                dither=self.dither,
+                energy_floor=0.0,
+                window_type=self.window,
+                sample_frequency=self.fs,
+                snip_edges=self.snip_edges,
+            )
+
+            if self.lfr_m != 1 or self.lfr_n != 1:
+                mat = apply_lfr(mat, self.lfr_m, self.lfr_n)
+            if self.cmvn is not None:
+                mat = apply_cmvn(mat, self.cmvn)
+            feat_length = mat.size(0)
+            feats.append(mat)
+            feats_lens.append(feat_length)
+
+        feats_lens = torch.as_tensor(feats_lens)
+        if batch_size == 1:
+            feats_pad = feats[0][None, :, :]
+        else:
+            feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0)
+        return feats_pad, feats_lens
+
+    def forward_fbank(
+        self, input: torch.Tensor, input_lengths: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward fbank.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+            """
+        batch_size = input.size(0)
+        feats = []
+        feats_lens = []
+        for i in range(batch_size):
+            waveform_length = input_lengths[i]
+            waveform = input[i][:waveform_length]
+            waveform = waveform * (1 << 15)
+            waveform = waveform.unsqueeze(0)
+            mat = kaldi.fbank(
+                waveform,
+                num_mel_bins=self.n_mels,
+                frame_length=self.frame_length,
+                frame_shift=self.frame_shift,
+                dither=self.dither,
+                energy_floor=0.0,
+                window_type=self.window,
+                sample_frequency=self.fs,
+            )
+
+            feat_length = mat.size(0)
+            feats.append(mat)
+            feats_lens.append(feat_length)
+
+        feats_lens = torch.as_tensor(feats_lens)
+        feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0)
+        return feats_pad, feats_lens
+
+    def forward_lfr_cmvn(
+        self, input: torch.Tensor, input_lengths: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward lfr cmvn.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+            """
+        batch_size = input.size(0)
+        feats = []
+        feats_lens = []
+        for i in range(batch_size):
+            mat = input[i, : input_lengths[i], :]
+            if self.lfr_m != 1 or self.lfr_n != 1:
+                mat = apply_lfr(mat, self.lfr_m, self.lfr_n)
+            if self.cmvn is not None:
+                mat = apply_cmvn(mat, self.cmvn)
+            feat_length = mat.size(0)
+            feats.append(mat)
+            feats_lens.append(feat_length)
+
+        feats_lens = torch.as_tensor(feats_lens)
+        feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0)
+        return feats_pad, feats_lens
+
+
+@tables.register("frontend_classes", "WavFrontendOnline")
+class WavFrontendOnline(nn.Module):
+    """Conventional frontend structure for streaming ASR/VAD."""
+
+    def __init__(
+        self,
+        cmvn_file: str = None,
+        fs: int = 16000,
+        window: str = "hamming",
+        n_mels: int = 80,
+        frame_length: int = 25,
+        frame_shift: int = 10,
+        filter_length_min: int = -1,
+        filter_length_max: int = -1,
+        lfr_m: int = 1,
+        lfr_n: int = 1,
+        dither: float = 1.0,
+        snip_edges: bool = True,
+        upsacle_samples: bool = True,
+        **kwargs,
+    ):
+        """Initialize WavFrontendOnline.
+        
+            Args:
+                cmvn_file: TODO.
+                fs: TODO.
+                window: TODO.
+                n_mels: TODO.
+                frame_length: TODO.
+                frame_shift: TODO.
+                filter_length_min: TODO.
+                filter_length_max: TODO.
+                lfr_m: TODO.
+                lfr_n: TODO.
+                dither: TODO.
+                snip_edges: TODO.
+                upsacle_samples: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        self.fs = fs
+        self.window = window
+        self.n_mels = n_mels
+        self.frame_length = frame_length
+        self.frame_shift = frame_shift
+        self.frame_sample_length = int(self.frame_length * self.fs / 1000)
+        self.frame_shift_sample_length = int(self.frame_shift * self.fs / 1000)
+        self.filter_length_min = filter_length_min
+        self.filter_length_max = filter_length_max
+        self.lfr_m = lfr_m
+        self.lfr_n = lfr_n
+        self.cmvn_file = cmvn_file
+        self.dither = dither
+        self.snip_edges = snip_edges
+        self.upsacle_samples = upsacle_samples
+        # self.waveforms = None
+        # self.reserve_waveforms = None
+        # self.fbanks = None
+        # self.fbanks_lens = None
+        self.cmvn = None if self.cmvn_file is None else load_cmvn(self.cmvn_file)
+        # self.input_cache = None
+        # self.lfr_splice_cache = []
+
+    def output_size(self) -> int:
+        """Output size."""
+        return self.n_mels * self.lfr_m
+
+    @staticmethod
+    def apply_cmvn(inputs: torch.Tensor, cmvn: torch.Tensor) -> torch.Tensor:
+        """
+        Apply CMVN with mvn data
+        """
+
+        device = inputs.device
+        dtype = inputs.dtype
+        frame, dim = inputs.shape
+
+        means = np.tile(cmvn[0:1, :dim], (frame, 1))
+        vars = np.tile(cmvn[1:2, :dim], (frame, 1))
+        inputs += torch.from_numpy(means).type(dtype).to(device)
+        inputs *= torch.from_numpy(vars).type(dtype).to(device)
+
+        return inputs.type(torch.float32)
+
+    @staticmethod
+    def apply_lfr(
+        inputs: torch.Tensor, lfr_m: int, lfr_n: int, is_final: bool = False
+    ) -> Tuple[torch.Tensor, torch.Tensor, int]:
+        """
+        Apply lfr with data
+        """
+
+        LFR_inputs = []
+        # inputs = torch.vstack((inputs_lfr_cache, inputs))
+        T = inputs.shape[0]  # include the right context
+        T_lfr = int(
+            np.ceil((T - (lfr_m - 1) // 2) / lfr_n)
+        )  # minus the right context: (lfr_m - 1) // 2
+        splice_idx = T_lfr
+        feat_dim = inputs.shape[-1]
+        ori_inputs = inputs
+        strides = (lfr_n * feat_dim, 1)
+        sizes = (T_lfr, lfr_m * feat_dim)
+        last_idx = (T - lfr_m) // lfr_n + 1
+        num_padding = lfr_m - (T - last_idx * lfr_n)
+        if is_final:
+            if num_padding > 0:
+                num_padding = (2 * lfr_m - 2 * T + (T_lfr - 1 + last_idx) * lfr_n) / 2 * (T_lfr - last_idx)
+                inputs = torch.vstack([inputs] + [inputs[-1:]] * int(num_padding))
+        else:
+            if num_padding > 0:
+                sizes = (last_idx, lfr_m * feat_dim)
+                splice_idx = last_idx
+        splice_idx = min(T - 1, splice_idx * lfr_n)
+        LFR_outputs = inputs[:splice_idx].as_strided(sizes, strides)
+        lfr_splice_cache = ori_inputs[splice_idx:, :]
+        return LFR_outputs.clone().type(torch.float32), lfr_splice_cache, splice_idx
+
+    @staticmethod
+    def compute_frame_num(
+        sample_length: int, frame_sample_length: int, frame_shift_sample_length: int
+    ) -> int:
+        """Compute frame num.
+        
+            Args:
+                sample_length: TODO.
+                frame_sample_length: TODO.
+                frame_shift_sample_length: TODO.
+            """
+        frame_num = int((sample_length - frame_sample_length) / frame_shift_sample_length + 1)
+        return frame_num if frame_num >= 1 and sample_length >= frame_sample_length else 0
+
+    def forward_fbank(
+        self,
+        input: torch.Tensor,
+        input_lengths: torch.Tensor,
+        cache: dict = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Forward fbank.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+                cache: State cache dict for streaming inference.
+                **kwargs: Additional keyword arguments.
+            """
+        if cache is None:
+            cache = {}
+        batch_size = input.size(0)
+
+        input = torch.cat((cache["input_cache"], input), dim=1)
+        frame_num = self.compute_frame_num(
+            input.shape[-1], self.frame_sample_length, self.frame_shift_sample_length
+        )
+        # update self.in_cache
+        cache["input_cache"] = input[
+            :, -(input.shape[-1] - frame_num * self.frame_shift_sample_length) :
+        ]
+        waveforms = torch.empty(0)
+        feats_pad = torch.empty(0)
+        feats_lens = torch.empty(0)
+        if frame_num:
+            waveforms = []
+            feats = []
+            feats_lens = []
+            for i in range(batch_size):
+                waveform = input[i]
+                # we need accurate wave samples that used for fbank extracting
+                waveforms.append(
+                    waveform[
+                        : (
+                            (frame_num - 1) * self.frame_shift_sample_length
+                            + self.frame_sample_length
+                        )
+                    ]
+                )
+                waveform = waveform * (1 << 15)
+                waveform = waveform.unsqueeze(0)
+                mat = kaldi.fbank(
+                    waveform,
+                    num_mel_bins=self.n_mels,
+                    frame_length=self.frame_length,
+                    frame_shift=self.frame_shift,
+                    dither=self.dither,
+                    energy_floor=0.0,
+                    window_type=self.window,
+                    sample_frequency=self.fs,
+                )
+
+                feat_length = mat.size(0)
+                feats.append(mat)
+                feats_lens.append(feat_length)
+
+            waveforms = torch.stack(waveforms)
+            feats_lens = torch.as_tensor(feats_lens)
+            feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0)
+        cache["fbanks"] = feats_pad
+        cache["fbanks_lens"] = copy.deepcopy(feats_lens)
+        return waveforms, feats_pad, feats_lens
+
+    def forward_lfr_cmvn(
+        self,
+        input: torch.Tensor,
+        input_lengths: torch.Tensor,
+        is_final: bool = False,
+        cache: dict = None,
+        **kwargs,
+    ):
+        """Forward lfr cmvn.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+                is_final: Whether this is the final chunk in streaming.
+                cache: State cache dict for streaming inference.
+                **kwargs: Additional keyword arguments.
+            """
+        if cache is None:
+            cache = {}
+        batch_size = input.size(0)
+        feats = []
+        feats_lens = []
+        lfr_splice_frame_idxs = []
+        for i in range(batch_size):
+            mat = input[i, : input_lengths[i], :]
+            if self.lfr_m != 1 or self.lfr_n != 1:
+                # update self.lfr_splice_cache in self.apply_lfr
+                # mat, self.lfr_splice_cache[i], lfr_splice_frame_idx = self.apply_lfr(mat, self.lfr_m, self.lfr_n, self.lfr_splice_cache[i],
+                mat, cache["lfr_splice_cache"][i], lfr_splice_frame_idx = self.apply_lfr(
+                    mat, self.lfr_m, self.lfr_n, is_final
+                )
+            if self.cmvn_file is not None:
+                mat = self.apply_cmvn(mat, self.cmvn)
+            feat_length = mat.size(0)
+            feats.append(mat)
+            feats_lens.append(feat_length)
+            lfr_splice_frame_idxs.append(lfr_splice_frame_idx)
+
+        feats_lens = torch.as_tensor(feats_lens)
+        feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0)
+        lfr_splice_frame_idxs = torch.as_tensor(lfr_splice_frame_idxs)
+        return feats_pad, feats_lens, lfr_splice_frame_idxs
+
+    def forward(self, input: torch.Tensor, input_lengths: torch.Tensor, **kwargs):
+        """Forward pass for training.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+                **kwargs: Additional keyword arguments.
+            """
+        is_final = kwargs.get("is_final", False)
+        cache = kwargs.get("cache", {})
+        if len(cache) == 0:
+            self.init_cache(cache)
+
+        batch_size = input.shape[0]
+        assert (
+            batch_size == 1
+        ), "we support to extract feature online only when the batch size is equal to 1 now"
+
+        waveforms, feats, feats_lengths = self.forward_fbank(
+            input, input_lengths, cache=cache
+        )  # input shape: B T D
+
+        if feats.shape[0]:
+
+            cache["waveforms"] = torch.cat((cache["reserve_waveforms"], waveforms), dim=1)
+
+            if not cache["lfr_splice_cache"]:  # 初始化splice_cache
+                for i in range(batch_size):
+                    cache["lfr_splice_cache"].append(
+                        feats[i][0, :].unsqueeze(dim=0).repeat((self.lfr_m - 1) // 2, 1)
+                    )
+            # need the number of the input frames + self.lfr_splice_cache[0].shape[0] is greater than self.lfr_m
+            if feats_lengths[0] + cache["lfr_splice_cache"][0].shape[0] >= self.lfr_m:
+                lfr_splice_cache_tensor = torch.stack(cache["lfr_splice_cache"])  # B T D
+                feats = torch.cat((lfr_splice_cache_tensor, feats), dim=1)
+                feats_lengths += lfr_splice_cache_tensor[0].shape[0]
+                frame_from_waveforms = int(
+                    (cache["waveforms"].shape[1] - self.frame_sample_length)
+                    / self.frame_shift_sample_length
+                    + 1
+                )
+                minus_frame = (
+                    (self.lfr_m - 1) // 2 if cache["reserve_waveforms"].numel() == 0 else 0
+                )
+                feats, feats_lengths, lfr_splice_frame_idxs = self.forward_lfr_cmvn(
+                    feats, feats_lengths, is_final, cache=cache
+                )
+                if self.lfr_m == 1:
+                    cache["reserve_waveforms"] = torch.empty(0)
+                else:
+                    reserve_frame_idx = lfr_splice_frame_idxs[0] - minus_frame
+                    # print('reserve_frame_idx:  ' + str(reserve_frame_idx))
+                    # print('frame_frame:  ' + str(frame_from_waveforms))
+                    cache["reserve_waveforms"] = cache["waveforms"][
+                        :,
+                        reserve_frame_idx
+                        * self.frame_shift_sample_length : frame_from_waveforms
+                        * self.frame_shift_sample_length,
+                    ]
+                    sample_length = (
+                        frame_from_waveforms - 1
+                    ) * self.frame_shift_sample_length + self.frame_sample_length
+                    cache["waveforms"] = cache["waveforms"][:, :sample_length]
+            else:
+                # update self.reserve_waveforms and self.lfr_splice_cache
+                cache["reserve_waveforms"] = cache["waveforms"][
+                    :, : -(self.frame_sample_length - self.frame_shift_sample_length)
+                ]
+                for i in range(batch_size):
+                    cache["lfr_splice_cache"][i] = torch.cat(
+                        (cache["lfr_splice_cache"][i], feats[i]), dim=0
+                    )
+                return torch.empty(0), feats_lengths
+        else:
+            if is_final:
+                cache["waveforms"] = (
+                    waveforms
+                    if cache["reserve_waveforms"].numel() == 0
+                    else cache["reserve_waveforms"]
+                )
+                feats = torch.stack(cache["lfr_splice_cache"])
+                feats_lengths = torch.zeros(batch_size, dtype=torch.int) + feats.shape[1]
+                feats, feats_lengths, _ = self.forward_lfr_cmvn(
+                    feats, feats_lengths, is_final, cache=cache
+                )
+        # if is_final:
+        #     self.init_cache(cache)
+        return feats, feats_lengths
+
+    def init_cache(self, cache: dict = None):
+        """Init cache.
+        
+            Args:
+                cache: State cache dict for streaming inference.
+            """
+        if cache is None:
+            cache = {}
+        cache["reserve_waveforms"] = torch.empty(0)
+        cache["input_cache"] = torch.empty(0)
+        cache["lfr_splice_cache"] = []
+        cache["waveforms"] = None
+        cache["fbanks"] = None
+        cache["fbanks_lens"] = None
+        return cache
+
+
+class WavFrontendMel23(nn.Module):
+    """Conventional frontend structure for ASR."""
+
+    def __init__(
+        self,
+        fs: int = 16000,
+        frame_length: int = 25,
+        frame_shift: int = 10,
+        lfr_m: int = 1,
+        lfr_n: int = 1,
+        **kwargs,
+    ):
+        """Initialize WavFrontendMel23.
+        
+            Args:
+                fs: TODO.
+                frame_length: TODO.
+                frame_shift: TODO.
+                lfr_m: TODO.
+                lfr_n: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        self.fs = fs
+        self.frame_length = frame_length
+        self.frame_shift = frame_shift
+        self.lfr_m = lfr_m
+        self.lfr_n = lfr_n
+        self.n_mels = 23
+
+    def output_size(self) -> int:
+        """Output size."""
+        return self.n_mels * (2 * self.lfr_m + 1)
+
+    def forward(
+        self, input: torch.Tensor, input_lengths: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass for training.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+            """
+        batch_size = input.size(0)
+        feats = []
+        feats_lens = []
+        for i in range(batch_size):
+            waveform_length = input_lengths[i]
+            waveform = input[i][:waveform_length]
+            waveform = waveform.numpy()
+            mat = eend_ola_feature.stft(waveform, self.frame_length, self.frame_shift)
+            mat = eend_ola_feature.transform(mat)
+            mat = eend_ola_feature.splice(mat, context_size=self.lfr_m)
+            mat = mat[:: self.lfr_n]
+            mat = torch.from_numpy(mat)
+            feat_length = mat.size(0)
+            feats.append(mat)
+            feats_lens.append(feat_length)
+
+        feats_lens = torch.as_tensor(feats_lens)
+        feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0)
+        return feats_pad, feats_lens

funasr/frontends/whisper_frontend.py

@@ -0,0 +1,141 @@
+from typing import Tuple
+import torch
+import torch.nn as nn
+
+
+from funasr.register import tables
+from torch.nn.utils.rnn import pad_sequence
+
+
+@tables.register("frontend_classes", "WhisperFrontend")
+class WhisperFrontend(nn.Module):
+    """Speech Representation Using Encoder Outputs from OpenAI's Whisper Model:
+
+    URL: https://github.com/openai/whisper
+    """
+
+    def __init__(
+        self,
+        fs: int = 16000,
+        whisper_model: str = None,
+        do_pad_trim: bool = True,
+        n_mels: int = 80,
+        permute: bool = False,
+        **kwargs,
+    ):
+        """Initialize WhisperFrontend.
+        
+            Args:
+                fs: TODO.
+                whisper_model: Whisper Model instance.
+                do_pad_trim: TODO.
+                n_mels: TODO.
+                permute: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        assert fs == 16000
+        self.fs = fs
+        import whisper
+        from whisper.audio import HOP_LENGTH, N_FFT, N_SAMPLES
+
+        self.n_fft = N_FFT
+        self.win_length = N_FFT
+        self.hop_length = HOP_LENGTH
+        self.pad_samples = N_SAMPLES
+        self.frame_shift = int(self.hop_length / self.fs * 1000)
+        self.lfr_n = 1
+        self.n_mels = n_mels
+        if whisper_model == "large-v3" or whisper_model == "large":
+            self.n_mels = 128
+
+        filters_path = kwargs.get("filters_path", None)
+        self.filters_path = filters_path
+        if filters_path is not None:
+            from funasr.models.sense_voice.whisper_lib.audio import mel_filters
+
+            self.mel_filters = mel_filters
+        else:
+            self.mel_filters = whisper.audio.mel_filters
+        self.do_pad_trim = do_pad_trim
+        if do_pad_trim:
+            self.pad_or_trim = whisper.pad_or_trim
+        self.permute = permute
+
+        # assert whisper_model in whisper.available_models()
+
+    def output_size(self) -> int:
+        """Output size."""
+        return self.n_mels
+
+    def log_mel_spectrogram(
+        self,
+        audio: torch.Tensor,
+        ilens: torch.Tensor = None,
+    ) -> torch.Tensor:
+        """Log mel spectrogram.
+        
+            Args:
+                audio: TODO.
+                ilens: TODO.
+            """
+        window = torch.hann_window(self.win_length).to(audio.device)
+        stft = torch.stft(audio, self.n_fft, self.hop_length, window=window, return_complex=True)
+
+        # whisper deletes the last frame by default (Shih-Lun)
+        magnitudes = stft[..., :-1].abs() ** 2
+        if self.filters_path is not None:
+            filters = self.mel_filters(audio.device, self.n_mels, self.filters_path)
+        else:
+            filters = self.mel_filters(audio.device, self.n_mels)
+        mel_spec = filters @ magnitudes
+
+        log_spec = torch.clamp(mel_spec, min=1e-10).log10()
+
+        if ilens is not None:
+            olens = ilens // self.hop_length
+        else:
+            olens = None
+
+        log_spec = torch.maximum(
+            log_spec,
+            log_spec.view(audio.size(0), -1).max(dim=-1)[0][:, None, None] - 8.0,
+        )
+        log_spec = (log_spec + 4.0) / 4.0
+
+        return log_spec, olens
+
+    def forward(
+        self,
+        input: torch.Tensor,
+        input_lengths: torch.Tensor,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass for training.
+        
+            Args:
+                input: Input audio/text data.
+                input_lengths: Lengths of input.
+                **kwargs: Additional keyword arguments.
+            """
+        batch_size = input.size(0)
+        feats = []
+        feats_lens = []
+        input = input.to(torch.float32)
+        for i in range(batch_size):
+            if self.do_pad_trim:
+                feat = self.pad_or_trim(input[i], self.pad_samples)
+            else:
+                feat = input[i]
+            feat, feat_len = self.log_mel_spectrogram(feat[None, :], input_lengths[0])
+            feats.append(feat[0])
+            feats_lens.append(feat_len)
+        feats_lens = torch.as_tensor(feats_lens)
+
+        if batch_size == 1:
+            feats_pad = feats[0][None, :, :]
+        else:
+            feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0)
+        if self.permute:
+            feats_pad = feats_pad.permute(0, 2, 1)
+        return feats_pad, feats_lens

funasr/frontends/windowing.py

@@ -0,0 +1,74 @@
+#!/usr/bin/env python3
+#  2020, Technische Universität München;  Ludwig Kürzinger
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+"""Sliding Window for raw audio input data."""
+
+import torch
+import torch.nn as nn
+from typing import Tuple
+
+
+class SlidingWindow(nn.Module):
+    """Sliding Window.
+    Provides a sliding window over a batched continuous raw audio tensor.
+    Optionally, provides padding (Currently not implemented).
+    Combine this module with a pre-encoder compatible with raw audio data,
+    for example Sinc convolutions.
+    Known issues:
+    Output length is calculated incorrectly if audio shorter than win_length.
+    WARNING: trailing values are discarded - padding not implemented yet.
+    There is currently no additional window function applied to input values.
+    """
+
+    def __init__(
+        self,
+        win_length: int = 400,
+        hop_length: int = 160,
+        channels: int = 1,
+        padding: int = None,
+        fs=None,
+    ):
+        """Initialize.
+        Args:
+            win_length: Length of frame.
+            hop_length: Relative starting point of next frame.
+            channels: Number of input channels.
+            padding: Padding (placeholder, currently not implemented).
+            fs:  Sampling rate (placeholder for compatibility, not used).
+        """
+        super().__init__()
+        self.fs = fs
+        self.win_length = win_length
+        self.hop_length = hop_length
+        self.channels = channels
+        self.padding = padding
+
+    def forward(
+        self, input: torch.Tensor, input_lengths: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Apply a sliding window on the input.
+        Args:
+            input: Input (B, T, C*D) or (B, T*C*D), with D=C=1.
+            input_lengths: Input lengths within batch.
+        Returns:
+            Tensor: Output with dimensions (B, T, C, D), with D=win_length.
+            Tensor: Output lengths within batch.
+        """
+        input_size = input.size()
+        B = input_size[0]
+        T = input_size[1]
+        C = self.channels
+        D = self.win_length
+        # (B, T, C) --> (T, B, C)
+        continuous = input.view(B, T, C).permute(1, 0, 2)
+        windowed = continuous.unfold(0, D, self.hop_length)
+        # (T, B, C, D) --> (B, T, C, D)
+        output = windowed.permute(1, 0, 2, 3).contiguous()
+        # After unfold(), windowed lengths change:
+        output_lengths = (input_lengths - self.win_length) // self.hop_length + 1
+        return output, output_lengths
+
+    def output_size(self) -> int:
+        """Return output length of feature dimension D, i.e. the window length."""
+        return self.win_length

funasr/losses/cr_ctc.py

@@ -0,0 +1,55 @@
+"""Consistency-Regularized CTC (CR-CTC) loss.
+
+Based on: "Improving CTC-based Speech Recognition via Consistency Regularization"
+Key idea: Run encoder twice (with/without SpecAug), compute KL divergence between
+the two CTC outputs as a consistency regularization term.
+
+Usage in training:
+    cr_loss = cr_ctc_loss(ctc_logprobs_aug, ctc_logprobs_clean, input_lengths)
+    total_loss = ctc_loss + cr_loss_scale * cr_loss
+"""
+
+import torch
+import torch.nn.functional as F
+
+
+def cr_ctc_loss(
+    log_probs_aug: torch.Tensor,
+    log_probs_clean: torch.Tensor,
+    input_lengths: torch.Tensor,
+) -> torch.Tensor:
+    """Compute CR-CTC consistency regularization loss.
+
+    Computes symmetric KL divergence between augmented and clean encoder outputs.
+
+    Args:
+        log_probs_aug: CTC log probabilities from augmented input (B, T, V)
+        log_probs_clean: CTC log probabilities from clean input (B, T, V)
+        input_lengths: Valid lengths for each sample (B,)
+
+    Returns:
+        Scalar loss value (mean over batch and time).
+    """
+    batch_size, max_len, _ = log_probs_aug.shape
+
+    # Create mask for valid positions
+    mask = torch.arange(max_len, device=input_lengths.device)[None, :] < input_lengths[:, None]
+    mask = mask.unsqueeze(-1)  # (B, T, 1)
+
+    # Convert log probs to probs for KL computation
+    probs_aug = log_probs_aug.exp()
+    probs_clean = log_probs_clean.exp()
+
+    # Symmetric KL divergence: 0.5 * (KL(p||q) + KL(q||p))
+    # KL(p||q) = sum(p * (log_p - log_q))
+    kl_aug_to_clean = (probs_aug * (log_probs_aug - log_probs_clean)) * mask
+    kl_clean_to_aug = (probs_clean * (log_probs_clean - log_probs_aug)) * mask
+
+    # Mean over valid positions
+    num_valid = mask.sum()
+    if num_valid > 0:
+        loss = 0.5 * (kl_aug_to_clean.sum() + kl_clean_to_aug.sum()) / num_valid
+    else:
+        loss = torch.tensor(0.0, device=log_probs_aug.device)
+
+    return loss

funasr/losses/label_smoothing_loss.py

@@ -0,0 +1,136 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+# Copyright 2019 Shigeki Karita
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+"""Label smoothing module."""
+
+import torch
+from torch import nn
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+
+
+class LabelSmoothingLoss(nn.Module):
+    """Label-smoothing loss.
+
+    :param int size: the number of class
+    :param int padding_idx: ignored class id
+    :param float smoothing: smoothing rate (0.0 means the conventional CE)
+    :param bool normalize_length: normalize loss by sequence length if True
+    :param torch.nn.Module criterion: loss function to be smoothed
+    """
+
+    def __init__(
+        self,
+        size,
+        padding_idx,
+        smoothing,
+        normalize_length=False,
+        criterion=nn.KLDivLoss(reduction="none"),
+    ):
+        """Construct an LabelSmoothingLoss object."""
+        super(LabelSmoothingLoss, self).__init__()
+        self.criterion = criterion
+        self.padding_idx = padding_idx
+        self.confidence = 1.0 - smoothing
+        self.smoothing = smoothing
+        self.size = size
+        self.true_dist = None
+        self.normalize_length = normalize_length
+
+    def forward(self, x, target):
+        """Compute loss between x and target.
+
+        :param torch.Tensor x: prediction (batch, seqlen, class)
+        :param torch.Tensor target:
+            target signal masked with self.padding_id (batch, seqlen)
+        :return: scalar float value
+        :rtype torch.Tensor
+        """
+        assert x.size(2) == self.size
+        batch_size = x.size(0)
+        x = x.contiguous().view(-1, self.size)
+        target = target.contiguous().view(-1)
+        with torch.no_grad():
+            true_dist = x.clone()
+            true_dist.fill_(self.smoothing / (self.size - 1))
+            ignore = target == self.padding_idx  # (B,)
+            total = len(target) - ignore.sum().item()
+            target = target.masked_fill(ignore, 0)  # avoid -1 index
+            true_dist.scatter_(1, target.unsqueeze(1), self.confidence)
+        kl = self.criterion(torch.log_softmax(x, dim=1), true_dist)
+        denom = total if self.normalize_length else batch_size
+        return kl.masked_fill(ignore.unsqueeze(1), 0).sum() / denom
+
+
+class SequenceBinaryCrossEntropy(nn.Module):
+    def __init__(self, normalize_length=False, criterion=nn.BCEWithLogitsLoss(reduction="none")):
+        """Initialize SequenceBinaryCrossEntropy.
+        
+            Args:
+                normalize_length: TODO.
+                criterion: TODO.
+            """
+        super().__init__()
+        self.normalize_length = normalize_length
+        self.criterion = criterion
+
+    def forward(self, pred, label, lengths):
+        """Forward pass for training.
+        
+            Args:
+                pred: TODO.
+                label: TODO.
+                lengths: TODO.
+            """
+        pad_mask = make_pad_mask(lengths, maxlen=pred.shape[1]).to(pred.device)
+        loss = self.criterion(pred, label)
+        denom = (~pad_mask).sum() if self.normalize_length else pred.shape[0]
+        return loss.masked_fill(pad_mask.unsqueeze(-1), 0).sum() / denom
+
+
+class NllLoss(nn.Module):
+    """Nll loss.
+
+    :param int size: the number of class
+    :param int padding_idx: ignored class id
+    :param bool normalize_length: normalize loss by sequence length if True
+    :param torch.nn.Module criterion: loss function
+    """
+
+    def __init__(
+        self,
+        size,
+        padding_idx,
+        normalize_length=False,
+        criterion=nn.NLLLoss(reduction="none"),
+    ):
+        """Construct an NllLoss object."""
+        super(NllLoss, self).__init__()
+        self.criterion = criterion
+        self.padding_idx = padding_idx
+        self.size = size
+        self.true_dist = None
+        self.normalize_length = normalize_length
+
+    def forward(self, x, target):
+        """Compute loss between x and target.
+
+        :param torch.Tensor x: prediction (batch, seqlen, class)
+        :param torch.Tensor target:
+            target signal masked with self.padding_id (batch, seqlen)
+        :return: scalar float value
+        :rtype torch.Tensor
+        """
+        assert x.size(2) == self.size
+        batch_size = x.size(0)
+        x = x.view(-1, self.size)
+        target = target.view(-1)
+        with torch.no_grad():
+            ignore = target == self.padding_idx  # (B,)
+            total = len(target) - ignore.sum().item()
+            target = target.masked_fill(ignore, 0)  # avoid -1 index
+        kl = self.criterion(x, target)
+        denom = total if self.normalize_length else batch_size
+        return kl.masked_fill(ignore, 0).sum() / denom

funasr/metrics/common.py

@@ -0,0 +1,245 @@
+#!/usr/bin/env python3
+# encoding: utf-8
+
+# Copyright 2017 Johns Hopkins University (Shinji Watanabe)
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+"""Common functions for ASR."""
+
+import json
+import logging
+import sys
+from itertools import groupby
+
+from rapidfuzz.distance import Levenshtein
+import numpy as np
+import six
+
+
+def end_detect(ended_hyps, i, M=3, D_end=np.log(1 * np.exp(-10))):
+    """End detection.
+
+    described in Eq. (50) of S. Watanabe et al
+    "Hybrid CTC/Attention Architecture for End-to-End Speech Recognition"
+
+    :param ended_hyps:
+    :param i:
+    :param M:
+    :param D_end:
+    :return:
+    """
+    if len(ended_hyps) == 0:
+        return False
+    count = 0
+    best_hyp = sorted(ended_hyps, key=lambda x: x["score"], reverse=True)[0]
+    for m in six.moves.range(M):
+        # get ended_hyps with their length is i - m
+        hyp_length = i - m
+        hyps_same_length = [x for x in ended_hyps if len(x["yseq"]) == hyp_length]
+        if len(hyps_same_length) > 0:
+            best_hyp_same_length = sorted(hyps_same_length, key=lambda x: x["score"], reverse=True)[
+                0
+            ]
+            if best_hyp_same_length["score"] - best_hyp["score"] < D_end:
+                count += 1
+
+    if count == M:
+        return True
+    else:
+        return False
+
+
+# TODO(takaaki-hori): add different smoothing methods
+def label_smoothing_dist(odim, lsm_type, transcript=None, blank=0):
+    """Obtain label distribution for loss smoothing.
+
+    :param odim:
+    :param lsm_type:
+    :param blank:
+    :param transcript:
+    :return:
+    """
+    if transcript is not None:
+        with open(transcript, "rb") as f:
+            trans_json = json.load(f)["utts"]
+
+    if lsm_type == "unigram":
+        assert transcript is not None, "transcript is required for %s label smoothing" % lsm_type
+        labelcount = np.zeros(odim)
+        for k, v in trans_json.items():
+            ids = np.array([int(n) for n in v["output"][0]["tokenid"].split()])
+            # to avoid an error when there is no text in an uttrance
+            if len(ids) > 0:
+                labelcount[ids] += 1
+        labelcount[odim - 1] = len(transcript)  # count <eos>
+        labelcount[labelcount == 0] = 1  # flooring
+        labelcount[blank] = 0  # remove counts for blank
+        labeldist = labelcount.astype(np.float32) / np.sum(labelcount)
+    else:
+        logging.error("Error: unexpected label smoothing type: %s" % lsm_type)
+        sys.exit()
+
+    return labeldist
+
+
+def get_vgg2l_odim(idim, in_channel=3, out_channel=128):
+    """Return the output size of the VGG frontend.
+
+    :param in_channel: input channel size
+    :param out_channel: output channel size
+    :return: output size
+    :rtype int
+    """
+    idim = idim / in_channel
+    idim = np.ceil(np.array(idim, dtype=np.float32) / 2)  # 1st max pooling
+    idim = np.ceil(np.array(idim, dtype=np.float32) / 2)  # 2nd max pooling
+    return int(idim) * out_channel  # numer of channels
+
+
+class ErrorCalculator(object):
+    """Calculate CER and WER for E2E_ASR and CTC models during training.
+
+    :param y_hats: numpy array with predicted text
+    :param y_pads: numpy array with true (target) text
+    :param char_list:
+    :param sym_space:
+    :param sym_blank:
+    :return:
+    """
+
+    def __init__(self, char_list, sym_space, sym_blank, report_cer=False, report_wer=False):
+        """Construct an ErrorCalculator object."""
+        super(ErrorCalculator, self).__init__()
+
+        self.report_cer = report_cer
+        self.report_wer = report_wer
+
+        self.char_list = char_list
+        self.space = sym_space
+        self.blank = sym_blank
+        self.idx_blank = self.char_list.index(self.blank)
+        if self.space in self.char_list:
+            self.idx_space = self.char_list.index(self.space)
+        else:
+            self.idx_space = None
+
+    def __call__(self, ys_hat, ys_pad, is_ctc=False):
+        """Calculate sentence-level WER/CER score.
+
+        :param torch.Tensor ys_hat: prediction (batch, seqlen)
+        :param torch.Tensor ys_pad: reference (batch, seqlen)
+        :param bool is_ctc: calculate CER score for CTC
+        :return: sentence-level WER score
+        :rtype float
+        :return: sentence-level CER score
+        :rtype float
+        """
+        cer, wer = None, None
+        if is_ctc:
+            return self.calculate_cer_ctc(ys_hat, ys_pad)
+        elif not self.report_cer and not self.report_wer:
+            return cer, wer
+
+        seqs_hat, seqs_true = self.convert_to_char(ys_hat, ys_pad)
+        if self.report_cer:
+            cer = self.calculate_cer(seqs_hat, seqs_true)
+
+        if self.report_wer:
+            wer = self.calculate_wer(seqs_hat, seqs_true)
+        return cer, wer
+
+    def calculate_cer_ctc(self, ys_hat, ys_pad):
+        """Calculate sentence-level CER score for CTC.
+
+        :param torch.Tensor ys_hat: prediction (batch, seqlen)
+        :param torch.Tensor ys_pad: reference (batch, seqlen)
+        :return: average sentence-level CER score
+        :rtype float
+        """
+
+        cers, char_ref_lens = [], []
+        for i, y in enumerate(ys_hat):
+            y_hat = [x[0] for x in groupby(y)]
+            y_true = ys_pad[i]
+            seq_hat, seq_true = [], []
+            for idx in y_hat:
+                idx = int(idx)
+                if idx != -1 and idx != self.idx_blank and idx != self.idx_space:
+                    seq_hat.append(self.char_list[int(idx)])
+
+            for idx in y_true:
+                idx = int(idx)
+                if idx != -1 and idx != self.idx_blank and idx != self.idx_space:
+                    seq_true.append(self.char_list[int(idx)])
+
+            hyp_chars = "".join(seq_hat)
+            ref_chars = "".join(seq_true)
+            if len(ref_chars) > 0:
+                cers.append(Levenshtein.distance(hyp_chars, ref_chars))
+                char_ref_lens.append(len(ref_chars))
+
+        cer_ctc = float(sum(cers)) / sum(char_ref_lens) if cers else None
+        return cer_ctc
+
+    def convert_to_char(self, ys_hat, ys_pad):
+        """Convert index to character.
+
+        :param torch.Tensor seqs_hat: prediction (batch, seqlen)
+        :param torch.Tensor seqs_true: reference (batch, seqlen)
+        :return: token list of prediction
+        :rtype list
+        :return: token list of reference
+        :rtype list
+        """
+        seqs_hat, seqs_true = [], []
+        for i, y_hat in enumerate(ys_hat):
+            y_true = ys_pad[i]
+            eos_true = np.where(y_true == -1)[0]
+            ymax = eos_true[0] if len(eos_true) > 0 else len(y_true)
+            # NOTE: padding index (-1) in y_true is used to pad y_hat
+            seq_hat = [self.char_list[int(idx)] for idx in y_hat[:ymax]]
+            seq_true = [self.char_list[int(idx)] for idx in y_true if int(idx) != -1]
+            seq_hat_text = "".join(seq_hat).replace(self.space, " ")
+            seq_hat_text = seq_hat_text.replace(self.blank, "")
+            seq_true_text = "".join(seq_true).replace(self.space, " ")
+            seqs_hat.append(seq_hat_text)
+            seqs_true.append(seq_true_text)
+        return seqs_hat, seqs_true
+
+    def calculate_cer(self, seqs_hat, seqs_true):
+        """Calculate sentence-level CER score.
+
+        :param list seqs_hat: prediction
+        :param list seqs_true: reference
+        :return: average sentence-level CER score
+        :rtype float
+        """
+
+        char_eds, char_ref_lens = [], []
+        for i, seq_hat_text in enumerate(seqs_hat):
+            seq_true_text = seqs_true[i]
+            hyp_chars = seq_hat_text.replace(" ", "")
+            ref_chars = seq_true_text.replace(" ", "")
+            char_eds.append(Levenshtein.distance(hyp_chars, ref_chars))
+            char_ref_lens.append(len(ref_chars))
+        ref_len = sum(char_ref_lens)
+        return float(sum(char_eds)) / ref_len if ref_len > 0 else None
+
+    def calculate_wer(self, seqs_hat, seqs_true):
+        """Calculate sentence-level WER score.
+
+        :param list seqs_hat: prediction
+        :param list seqs_true: reference
+        :return: average sentence-level WER score
+        :rtype float
+        """
+
+        word_eds, word_ref_lens = [], []
+        for i, seq_hat_text in enumerate(seqs_hat):
+            seq_true_text = seqs_true[i]
+            hyp_words = seq_hat_text.split()
+            ref_words = seq_true_text.split()
+            word_eds.append(Levenshtein.distance(hyp_words, ref_words))
+            word_ref_lens.append(len(ref_words))
+        ref_len = sum(word_ref_lens)
+        return float(sum(word_eds)) / ref_len if ref_len > 0 else None

funasr/metrics/compute_acc.py

@@ -0,0 +1,40 @@
+import torch
+
+
+def th_accuracy(pad_outputs, pad_targets, ignore_label):
+    """Calculate accuracy.
+
+    Args:
+        pad_outputs (Tensor): Prediction tensors (B * Lmax, D).
+        pad_targets (LongTensor): Target label tensors (B, Lmax, D).
+        ignore_label (int): Ignore label id.
+
+    Returns:
+        float: Accuracy value (0.0 - 1.0).
+
+    """
+    pad_pred = pad_outputs.view(
+        pad_targets.size(0), pad_targets.size(1), pad_outputs.size(1)
+    ).argmax(2)
+    mask = pad_targets != ignore_label
+    numerator = torch.sum(pad_pred.masked_select(mask) == pad_targets.masked_select(mask))
+    denominator = torch.sum(mask)
+    return float(numerator) / float(denominator)
+
+
+def compute_accuracy(pad_outputs, pad_targets, ignore_label):
+    """Calculate accuracy.
+
+    Args:
+        pad_outputs (LongTensor): Prediction tensors (B, Lmax).
+        pad_targets (LongTensor): Target label tensors (B, Lmax).
+        ignore_label (int): Ignore label id.
+
+    Returns:
+        float: Accuracy value (0.0 - 1.0).
+
+    """
+    mask = pad_targets != ignore_label
+    numerator = torch.sum(pad_outputs.masked_select(mask) == pad_targets.masked_select(mask))
+    denominator = torch.sum(mask)
+    return numerator.float() / denominator.float()  # (FIX:MZY):return torch.Tensor type

funasr/metrics/compute_eer.py

@@ -0,0 +1,66 @@
+import numpy as np
+from sklearn.metrics import roc_curve
+import argparse
+
+
+def _compute_eer(label, pred, positive_label=1):
+    """
+    Python compute equal error rate (eer)
+    ONLY tested on binary classification
+
+    :param label: ground-truth label, should be a 1-d list or np.array, each element represents the ground-truth label of one sample
+    :param pred: model prediction, should be a 1-d list or np.array, each element represents the model prediction of one sample
+    :param positive_label: the class that is viewed as positive class when computing EER
+    :return: equal error rate (EER)
+    """
+
+    # all fpr, tpr, fnr, fnr, threshold are lists (in the format of np.array)
+    fpr, tpr, threshold = roc_curve(label, pred, pos_label=positive_label)
+    fnr = 1 - tpr
+
+    # the threshold of fnr == fpr
+    eer_threshold = threshold[np.nanargmin(np.absolute((fnr - fpr)))]
+
+    # theoretically eer from fpr and eer from fnr should be identical but they can be slightly differ in reality
+    eer_1 = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+    eer_2 = fnr[np.nanargmin(np.absolute((fnr - fpr)))]
+
+    # return the mean of eer from fpr and from fnr
+    eer = (eer_1 + eer_2) / 2
+    return eer, eer_threshold
+
+
+def compute_eer(trials_path, scores_path):
+    """Compute eer.
+    
+        Args:
+            trials_path: TODO.
+            scores_path: TODO.
+        """
+    labels = []
+    for one_line in open(trials_path, "r"):
+        labels.append(one_line.strip().rsplit(" ", 1)[-1] == "target")
+    labels = np.array(labels, dtype=int)
+
+    scores = []
+    for one_line in open(scores_path, "r"):
+        scores.append(float(one_line.strip().rsplit(" ", 1)[-1]))
+    scores = np.array(scores, dtype=float)
+
+    eer, threshold = _compute_eer(labels, scores)
+    return eer, threshold
+
+
+def main():
+    """Main."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("trials", help="trial list")
+    parser.add_argument("scores", help="score file, normalized to [0, 1]")
+    args = parser.parse_args()
+
+    eer, threshold = compute_eer(args.trials, args.scores)
+    print("EER is {:.4f} at threshold {:.4f}".format(eer * 100.0, threshold))
+
+
+if __name__ == "__main__":
+    main()

funasr/metrics/compute_min_dcf.py

@@ -0,0 +1,208 @@
+#!/usr/bin/env python3
+# Copyright 2018  David Snyder
+# Apache 2.0
+
+# This script computes the minimum detection cost function, which is a common
+# error metric used in speaker recognition.  Compared to equal error-rate,
+# which assigns equal weight to false negatives and false positives, this
+# error-rate is usually used to assess performance in settings where achieving
+# a low false positive rate is more important than achieving a low false
+# negative rate.  See the NIST 2016 Speaker Recognition Evaluation Plan at
+# https://www.nist.gov/sites/default/files/documents/2016/10/07/sre16_eval_plan_v1.3.pdf
+# for more details about the metric.
+from __future__ import print_function
+from operator import itemgetter
+import sys, argparse, os
+
+
+def GetArgs():
+    """Getargs."""
+    parser = argparse.ArgumentParser(
+        description="Compute the minimum "
+        "detection cost function along with the threshold at which it occurs. "
+        "Usage: sid/compute_min_dcf.py [options...] <scores-file> "
+        "<trials-file> "
+        "E.g., sid/compute_min_dcf.py --p-target 0.01 --c-miss 1 --c-fa 1 "
+        "exp/scores/trials data/test/trials",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "--p-target",
+        type=float,
+        dest="p_target",
+        default=0.01,
+        help="The prior probability of the target speaker in a trial.",
+    )
+    parser.add_argument(
+        "--c-miss",
+        type=float,
+        dest="c_miss",
+        default=1,
+        help="Cost of a missed detection.  This is usually not changed.",
+    )
+    parser.add_argument(
+        "--c-fa",
+        type=float,
+        dest="c_fa",
+        default=1,
+        help="Cost of a spurious detection.  This is usually not changed.",
+    )
+    parser.add_argument(
+        "scores_filename",
+        help="Input scores file, with columns of the form " "<utt1> <utt2> <score>",
+    )
+    parser.add_argument(
+        "trials_filename",
+        help="Input trials file, with columns of the form " "<utt1> <utt2> <target/nontarget>",
+    )
+    sys.stderr.write(" ".join(sys.argv) + "\n")
+    args = parser.parse_args()
+    args = CheckArgs(args)
+    return args
+
+
+def CheckArgs(args):
+    """Checkargs.
+    
+        Args:
+            args: TODO.
+        """
+    if args.c_fa <= 0:
+        raise Exception("--c-fa must be greater than 0")
+    if args.c_miss <= 0:
+        raise Exception("--c-miss must be greater than 0")
+    if args.p_target <= 0 or args.p_target >= 1:
+        raise Exception("--p-target must be greater than 0 and less than 1")
+    return args
+
+
+# Creates a list of false-negative rates, a list of false-positive rates
+# and a list of decision thresholds that give those error-rates.
+def ComputeErrorRates(scores, labels):
+
+    # Sort the scores from smallest to largest, and also get the corresponding
+    # indexes of the sorted scores.  We will treat the sorted scores as the
+    # thresholds at which the the error-rates are evaluated.
+    """Computeerrorrates.
+    
+        Args:
+            scores: TODO.
+            labels: TODO.
+        """
+    sorted_indexes, thresholds = zip(
+        *sorted([(index, threshold) for index, threshold in enumerate(scores)], key=itemgetter(1))
+    )
+    labels = [labels[i] for i in sorted_indexes]
+    fns = []
+    tns = []
+
+    # At the end of this loop, fns[i] is the number of errors made by
+    # incorrectly rejecting scores less than thresholds[i]. And, tns[i]
+    # is the total number of times that we have correctly rejected scores
+    # less than thresholds[i].
+    for i in range(0, len(labels)):
+        if i == 0:
+            fns.append(labels[i])
+            tns.append(1 - labels[i])
+        else:
+            fns.append(fns[i - 1] + labels[i])
+            tns.append(tns[i - 1] + 1 - labels[i])
+    positives = sum(labels)
+    negatives = len(labels) - positives
+
+    # Now divide the false negatives by the total number of
+    # positives to obtain the false negative rates across
+    # all thresholds
+    fnrs = [fn / float(positives) for fn in fns]
+
+    # Divide the true negatives by the total number of
+    # negatives to get the true negative rate. Subtract these
+    # quantities from 1 to get the false positive rates.
+    fprs = [1 - tn / float(negatives) for tn in tns]
+    return fnrs, fprs, thresholds
+
+
+# Computes the minimum of the detection cost function.  The comments refer to
+# equations in Section 3 of the NIST 2016 Speaker Recognition Evaluation Plan.
+def ComputeMinDcf(fnrs, fprs, thresholds, p_target, c_miss, c_fa):
+    """Computemindcf.
+    
+        Args:
+            fnrs: TODO.
+            fprs: TODO.
+            thresholds: TODO.
+            p_target: TODO.
+            c_miss: TODO.
+            c_fa: TODO.
+        """
+    min_c_det = float("inf")
+    min_c_det_threshold = thresholds[0]
+    for i in range(0, len(fnrs)):
+        # See Equation (2).  it is a weighted sum of false negative
+        # and false positive errors.
+        c_det = c_miss * fnrs[i] * p_target + c_fa * fprs[i] * (1 - p_target)
+        if c_det < min_c_det:
+            min_c_det = c_det
+            min_c_det_threshold = thresholds[i]
+    # See Equations (3) and (4).  Now we normalize the cost.
+    c_def = min(c_miss * p_target, c_fa * (1 - p_target))
+    min_dcf = min_c_det / c_def
+    return min_dcf, min_c_det_threshold
+
+
+def compute_min_dcf(scores_filename, trials_filename, c_miss=1, c_fa=1, p_target=0.01):
+    """Compute min dcf.
+    
+        Args:
+            scores_filename: TODO.
+            trials_filename: TODO.
+            c_miss: TODO.
+            c_fa: TODO.
+            p_target: TODO.
+        """
+    scores_file = open(scores_filename, "r").readlines()
+    trials_file = open(trials_filename, "r").readlines()
+    c_miss = c_miss
+    c_fa = c_fa
+    p_target = p_target
+
+    scores = []
+    labels = []
+
+    trials = {}
+    for line in trials_file:
+        utt1, utt2, target = line.rstrip().split()
+        trial = utt1 + " " + utt2
+        trials[trial] = target
+
+    for line in scores_file:
+        utt1, utt2, score = line.rstrip().split()
+        trial = utt1 + " " + utt2
+        if trial in trials:
+            scores.append(float(score))
+            if trials[trial] == "target":
+                labels.append(1)
+            else:
+                labels.append(0)
+        else:
+            raise Exception("Missing entry for " + utt1 + " and " + utt2 + " " + scores_filename)
+
+    fnrs, fprs, thresholds = ComputeErrorRates(scores, labels)
+    mindcf, threshold = ComputeMinDcf(fnrs, fprs, thresholds, p_target, c_miss, c_fa)
+    return mindcf, threshold
+
+
+def main():
+    """Main."""
+    args = GetArgs()
+    mindcf, threshold = compute_min_dcf(
+        args.scores_filename, args.trials_filename, args.c_miss, args.c_fa, args.p_target
+    )
+    sys.stdout.write(
+        "minDCF is {0:.4f} at threshold {1:.4f} (p-target={2}, c-miss={3}, "
+        "c-fa={4})\n".format(mindcf, threshold, args.p_target, args.c_miss, args.c_fa)
+    )
+
+
+if __name__ == "__main__":
+    main()

funasr/metrics/wer.py

@@ -0,0 +1,249 @@
+import os
+import numpy as np
+import sys
+import hydra
+from omegaconf import DictConfig, OmegaConf, ListConfig
+
+
+def compute_wer(
+    ref_file,
+    hyp_file,
+    cer_file,
+    cn_postprocess=False,
+):
+    """Compute wer.
+    
+        Args:
+            ref_file: TODO.
+            hyp_file: TODO.
+            cer_file: TODO.
+            cn_postprocess: TODO.
+        """
+    rst = {
+        "Wrd": 0,
+        "Corr": 0,
+        "Ins": 0,
+        "Del": 0,
+        "Sub": 0,
+        "Snt": 0,
+        "Err": 0.0,
+        "S.Err": 0.0,
+        "wrong_words": 0,
+        "wrong_sentences": 0,
+    }
+
+    hyp_dict = {}
+    ref_dict = {}
+    with open(hyp_file, "r") as hyp_reader:
+        for line in hyp_reader:
+            key = line.strip().split()[0]
+            value = line.strip().split()[1:]
+            if cn_postprocess:
+                value = " ".join(value)
+                value = value.replace(" ", "")
+                # if value[0] == "请":
+                #     value = value[1:]
+                value = [x for x in value]
+            hyp_dict[key] = value
+    with open(ref_file, "r") as ref_reader:
+        for line in ref_reader:
+            key = line.strip().split()[0]
+            value = line.strip().split()[1:]
+            if cn_postprocess:
+                value = " ".join(value)
+                value = value.replace(" ", "")
+                value = [x for x in value]
+            ref_dict[key] = value
+
+    cer_detail_writer = open(cer_file, "w")
+    for hyp_key in hyp_dict:
+        if hyp_key in ref_dict:
+            out_item = compute_wer_by_line(hyp_dict[hyp_key], ref_dict[hyp_key])
+            rst["Wrd"] += out_item["nwords"]
+            rst["Corr"] += out_item["cor"]
+            rst["wrong_words"] += out_item["wrong"]
+            rst["Ins"] += out_item["ins"]
+            rst["Del"] += out_item["del"]
+            rst["Sub"] += out_item["sub"]
+            rst["Snt"] += 1
+            if out_item["wrong"] > 0:
+                rst["wrong_sentences"] += 1
+            cer_detail_writer.write(hyp_key + print_cer_detail(out_item) + "\n")
+            cer_detail_writer.write(
+                "ref:" + "\t" + " ".join(list(map(lambda x: x.lower(), ref_dict[hyp_key]))) + "\n"
+            )
+            cer_detail_writer.write(
+                "hyp:" + "\t" + " ".join(list(map(lambda x: x.lower(), hyp_dict[hyp_key]))) + "\n"
+            )
+            cer_detail_writer.flush()
+
+    if rst["Wrd"] > 0:
+        rst["Err"] = round(rst["wrong_words"] * 100 / rst["Wrd"], 2)
+    if rst["Snt"] > 0:
+        rst["S.Err"] = round(rst["wrong_sentences"] * 100 / rst["Snt"], 2)
+
+    cer_detail_writer.write("\n")
+    cer_detail_writer.write(
+        "%WER "
+        + str(rst["Err"])
+        + " [ "
+        + str(rst["wrong_words"])
+        + " / "
+        + str(rst["Wrd"])
+        + ", "
+        + str(rst["Ins"])
+        + " ins, "
+        + str(rst["Del"])
+        + " del, "
+        + str(rst["Sub"])
+        + " sub ]"
+        + "\n"
+    )
+    cer_detail_writer.write(
+        "%SER "
+        + str(rst["S.Err"])
+        + " [ "
+        + str(rst["wrong_sentences"])
+        + " / "
+        + str(rst["Snt"])
+        + " ]"
+        + "\n"
+    )
+    cer_detail_writer.write(
+        "Scored "
+        + str(len(hyp_dict))
+        + " sentences, "
+        + str(len(hyp_dict) - rst["Snt"])
+        + " not present in hyp."
+        + "\n"
+    )
+
+    cer_detail_writer.close()
+
+
+def compute_wer_by_line(hyp, ref):
+    """Compute wer by line.
+    
+        Args:
+            hyp: TODO.
+            ref: TODO.
+        """
+    hyp = list(map(lambda x: x.lower(), hyp))
+    ref = list(map(lambda x: x.lower(), ref))
+
+    len_hyp = len(hyp)
+    len_ref = len(ref)
+
+    cost_matrix = np.zeros((len_hyp + 1, len_ref + 1), dtype=np.int16)
+
+    ops_matrix = np.zeros((len_hyp + 1, len_ref + 1), dtype=np.int8)
+
+    for i in range(len_hyp + 1):
+        cost_matrix[i][0] = i
+    for j in range(len_ref + 1):
+        cost_matrix[0][j] = j
+
+    for i in range(1, len_hyp + 1):
+        for j in range(1, len_ref + 1):
+            if hyp[i - 1] == ref[j - 1]:
+                cost_matrix[i][j] = cost_matrix[i - 1][j - 1]
+            else:
+                substitution = cost_matrix[i - 1][j - 1] + 1
+                insertion = cost_matrix[i - 1][j] + 1
+                deletion = cost_matrix[i][j - 1] + 1
+
+                compare_val = [substitution, insertion, deletion]
+
+                min_val = min(compare_val)
+                operation_idx = compare_val.index(min_val) + 1
+                cost_matrix[i][j] = min_val
+                ops_matrix[i][j] = operation_idx
+
+    match_idx = []
+    i = len_hyp
+    j = len_ref
+    rst = {"nwords": len_ref, "cor": 0, "wrong": 0, "ins": 0, "del": 0, "sub": 0}
+    while i >= 0 or j >= 0:
+        i_idx = max(0, i)
+        j_idx = max(0, j)
+
+        if ops_matrix[i_idx][j_idx] == 0:  # correct
+            if i - 1 >= 0 and j - 1 >= 0:
+                match_idx.append((j - 1, i - 1))
+                rst["cor"] += 1
+
+            i -= 1
+            j -= 1
+
+        elif ops_matrix[i_idx][j_idx] == 2:  # insert
+            i -= 1
+            rst["ins"] += 1
+
+        elif ops_matrix[i_idx][j_idx] == 3:  # delete
+            j -= 1
+            rst["del"] += 1
+
+        elif ops_matrix[i_idx][j_idx] == 1:  # substitute
+            i -= 1
+            j -= 1
+            rst["sub"] += 1
+
+        if i < 0 and j >= 0:
+            rst["del"] += 1
+        elif j < 0 and i >= 0:
+            rst["ins"] += 1
+
+    match_idx.reverse()
+    wrong_cnt = cost_matrix[len_hyp][len_ref]
+    rst["wrong"] = wrong_cnt
+
+    return rst
+
+
+def print_cer_detail(rst):
+    """Print cer detail.
+    
+        Args:
+            rst: TODO.
+        """
+    return (
+        "("
+        + "nwords="
+        + str(rst["nwords"])
+        + ",cor="
+        + str(rst["cor"])
+        + ",ins="
+        + str(rst["ins"])
+        + ",del="
+        + str(rst["del"])
+        + ",sub="
+        + str(rst["sub"])
+        + ") corr:"
+        + "{:.2%}".format(rst["cor"] / rst["nwords"])
+        + ",cer:"
+        + "{:.2%}".format(rst["wrong"] / rst["nwords"])
+    )
+
+
+@hydra.main(config_name=None, version_base=None)
+def main_hydra(cfg: DictConfig):
+    """Main hydra.
+    
+        Args:
+            cfg: Configuration overrides.
+        """
+    ref_file = cfg.get("ref_file", None)
+    hyp_file = cfg.get("hyp_file", None)
+    cer_file = cfg.get("cer_file", None)
+    cn_postprocess = cfg.get("cn_postprocess", False)
+    if ref_file is None or hyp_file is None or cer_file is None:
+        print(
+            "usage : python -m  funasr.metrics.wer ++ref_file=test.ref ++hyp_file=test.hyp ++cer_file=test.wer ++cn_postprocess=false"
+        )
+        sys.exit(0)
+
+    compute_wer(ref_file, hyp_file, cer_file, cn_postprocess)
+
+
+if __name__ == "__main__":
+    main_hydra()

funasr/models/bat/model.py

@@ -0,0 +1,48 @@
+#!/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 logging
+from contextlib import contextmanager
+from typing import Dict, Optional, Tuple
+from distutils.version import LooseVersion
+
+from funasr.register import tables
+from funasr.utils import postprocess_utils
+from funasr.utils.datadir_writer import DatadirWriter
+from funasr.models.transducer.model import Transducer
+from funasr.train_utils.device_funcs import force_gatherable
+from funasr.models.transformer.scorers.ctc import CTCPrefixScorer
+from funasr.losses.label_smoothing_loss import LabelSmoothingLoss
+from funasr.models.transformer.scorers.length_bonus import LengthBonus
+from funasr.models.transformer.utils.nets_utils import get_transducer_task_io
+from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
+from funasr.models.transducer.beam_search_transducer import BeamSearchTransducer
+
+
+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):
+        """Autocast.
+        
+            Args:
+                enabled: TODO.
+            """
+        yield
+
+
+@tables.register("model_classes", "BAT")  # TODO: BAT training
+class BAT(Transducer):
+    """BAT (Boundary-Aware Transducer): Low-latency RNN-T model with boundary detection.
+
+    Inherits from Transducer. Designed for streaming ASR with reduced latency
+    by predicting token boundaries explicitly.
+    """
+
+    pass

funasr/models/bicif_paraformer/cif_predictor.py

@@ -0,0 +1,659 @@
+#!/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 torch
+
+from funasr.register import tables
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+
+
+class mae_loss(torch.nn.Module):
+
+    def __init__(self, normalize_length=False):
+        """Initialize mae_loss.
+        
+            Args:
+                normalize_length: TODO.
+            """
+        super(mae_loss, self).__init__()
+        self.normalize_length = normalize_length
+        self.criterion = torch.nn.L1Loss(reduction="sum")
+
+    def forward(self, token_length, pre_token_length):
+        """Forward pass for training.
+        
+            Args:
+                token_length: TODO.
+                pre_token_length: TODO.
+            """
+        loss_token_normalizer = token_length.size(0)
+        if self.normalize_length:
+            loss_token_normalizer = token_length.sum().type(torch.float32)
+        loss = self.criterion(token_length, pre_token_length)
+        loss = loss / loss_token_normalizer
+        return loss
+
+
+def cif(hidden, alphas, threshold):
+    """Cif.
+    
+        Args:
+            hidden: TODO.
+            alphas: TODO.
+            threshold: TODO.
+        """
+    batch_size, len_time, hidden_size = hidden.size()
+
+    # loop varss
+    integrate = torch.zeros([batch_size], device=hidden.device)
+    frame = torch.zeros([batch_size, hidden_size], device=hidden.device)
+    # intermediate vars along time
+    list_fires = []
+    list_frames = []
+
+    for t in range(len_time):
+        alpha = alphas[:, t]
+        distribution_completion = torch.ones([batch_size], device=hidden.device) - integrate
+
+        integrate += alpha
+        list_fires.append(integrate)
+
+        fire_place = integrate >= threshold
+        integrate = torch.where(
+            fire_place, integrate - torch.ones([batch_size], device=hidden.device), integrate
+        )
+        cur = torch.where(fire_place, distribution_completion, alpha)
+        remainds = alpha - cur
+
+        frame += cur[:, None] * hidden[:, t, :]
+        list_frames.append(frame)
+        frame = torch.where(
+            fire_place[:, None].repeat(1, hidden_size), remainds[:, None] * hidden[:, t, :], frame
+        )
+
+    fires = torch.stack(list_fires, 1)
+    frames = torch.stack(list_frames, 1)
+    list_ls = []
+    len_labels = torch.round(alphas.sum(-1)).int()
+    max_label_len = len_labels.max()
+    for b in range(batch_size):
+        fire = fires[b, :]
+        l = torch.index_select(frames[b, :, :], 0, torch.nonzero(fire >= threshold).squeeze(-1))
+        pad_l = torch.zeros([max_label_len - l.size(0), hidden_size], device=hidden.device)
+        list_ls.append(torch.cat([l, pad_l], 0))
+    return torch.stack(list_ls, 0), fires
+
+
+def cif_wo_hidden(alphas, threshold):
+    """Cif wo hidden.
+    
+        Args:
+            alphas: TODO.
+            threshold: TODO.
+        """
+    batch_size, len_time = alphas.size()
+
+    # loop varss
+    integrate = torch.zeros([batch_size], device=alphas.device)
+    # intermediate vars along time
+    list_fires = []
+
+    for t in range(len_time):
+        alpha = alphas[:, t]
+
+        integrate += alpha
+        list_fires.append(integrate)
+
+        fire_place = integrate >= threshold
+        integrate = torch.where(
+            fire_place,
+            integrate - torch.ones([batch_size], device=alphas.device) * threshold,
+            integrate,
+        )
+
+    fires = torch.stack(list_fires, 1)
+    return fires
+
+
+@tables.register("predictor_classes", "CifPredictorV3")
+class CifPredictorV3(torch.nn.Module):
+    def __init__(
+        self,
+        idim,
+        l_order,
+        r_order,
+        threshold=1.0,
+        dropout=0.1,
+        smooth_factor=1.0,
+        noise_threshold=0,
+        tail_threshold=0.0,
+        tf2torch_tensor_name_prefix_torch="predictor",
+        tf2torch_tensor_name_prefix_tf="seq2seq/cif",
+        smooth_factor2=1.0,
+        noise_threshold2=0,
+        upsample_times=5,
+        upsample_type="cnn",
+        use_cif1_cnn=True,
+        tail_mask=True,
+    ):
+        """Initialize CifPredictorV3.
+        
+            Args:
+                idim: TODO.
+                l_order: TODO.
+                r_order: TODO.
+                threshold: TODO.
+                dropout: TODO.
+                smooth_factor: TODO.
+                noise_threshold: TODO.
+                tail_threshold: TODO.
+                tf2torch_tensor_name_prefix_torch: TODO.
+                tf2torch_tensor_name_prefix_tf: TODO.
+                smooth_factor2: TODO.
+                noise_threshold2: TODO.
+                upsample_times: TODO.
+                upsample_type: TODO.
+                use_cif1_cnn: TODO.
+                tail_mask: TODO.
+            """
+        super(CifPredictorV3, self).__init__()
+
+        self.pad = torch.nn.ConstantPad1d((l_order, r_order), 0)
+        self.cif_conv1d = torch.nn.Conv1d(idim, idim, l_order + r_order + 1)
+        self.cif_output = torch.nn.Linear(idim, 1)
+        self.dropout = torch.nn.Dropout(p=dropout)
+        self.threshold = threshold
+        self.smooth_factor = smooth_factor
+        self.noise_threshold = noise_threshold
+        self.tail_threshold = tail_threshold
+        self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch
+        self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf
+
+        self.upsample_times = upsample_times
+        self.upsample_type = upsample_type
+        self.use_cif1_cnn = use_cif1_cnn
+        if self.upsample_type == "cnn":
+            self.upsample_cnn = torch.nn.ConvTranspose1d(
+                idim, idim, self.upsample_times, self.upsample_times
+            )
+            self.cif_output2 = torch.nn.Linear(idim, 1)
+        elif self.upsample_type == "cnn_blstm":
+            self.upsample_cnn = torch.nn.ConvTranspose1d(
+                idim, idim, self.upsample_times, self.upsample_times
+            )
+            self.blstm = torch.nn.LSTM(
+                idim, idim, 1, bias=True, batch_first=True, dropout=0.0, bidirectional=True
+            )
+            self.cif_output2 = torch.nn.Linear(idim * 2, 1)
+        elif self.upsample_type == "cnn_attn":
+            self.upsample_cnn = torch.nn.ConvTranspose1d(
+                idim, idim, self.upsample_times, self.upsample_times
+            )
+            from funasr.models.transformer.encoder import EncoderLayer as TransformerEncoderLayer
+            from funasr.models.transformer.attention import MultiHeadedAttention
+            from funasr.models.transformer.positionwise_feed_forward import PositionwiseFeedForward
+
+            positionwise_layer_args = (
+                idim,
+                idim * 2,
+                0.1,
+            )
+            self.self_attn = TransformerEncoderLayer(
+                idim,
+                MultiHeadedAttention(4, idim, 0.1),
+                PositionwiseFeedForward(*positionwise_layer_args),
+                0.1,
+                True,  # normalize_before,
+                False,  # concat_after,
+            )
+            self.cif_output2 = torch.nn.Linear(idim, 1)
+        self.smooth_factor2 = smooth_factor2
+        self.noise_threshold2 = noise_threshold2
+
+    def forward(
+        self,
+        hidden,
+        target_label=None,
+        mask=None,
+        ignore_id=-1,
+        mask_chunk_predictor=None,
+        target_label_length=None,
+    ):
+        """Forward pass for training.
+        
+            Args:
+                hidden: TODO.
+                target_label: TODO.
+                mask: TODO.
+                ignore_id: TODO.
+                mask_chunk_predictor: TODO.
+                target_label_length: TODO.
+            """
+        h = hidden
+        context = h.transpose(1, 2)
+        queries = self.pad(context)
+        output = torch.relu(self.cif_conv1d(queries))
+
+        # alphas2 is an extra head for timestamp prediction
+        if not self.use_cif1_cnn:
+            _output = context
+        else:
+            _output = output
+        if self.upsample_type == "cnn":
+            output2 = self.upsample_cnn(_output)
+            output2 = output2.transpose(1, 2)
+        elif self.upsample_type == "cnn_blstm":
+            output2 = self.upsample_cnn(_output)
+            output2 = output2.transpose(1, 2)
+            output2, (_, _) = self.blstm(output2)
+        elif self.upsample_type == "cnn_attn":
+            output2 = self.upsample_cnn(_output)
+            output2 = output2.transpose(1, 2)
+            output2, _ = self.self_attn(output2, mask)
+        
+        alphas2 = torch.sigmoid(self.cif_output2(output2))
+        alphas2 = torch.nn.functional.relu(alphas2 * self.smooth_factor2 - self.noise_threshold2)
+        # repeat the mask in T demension to match the upsampled length
+        if mask is not None:
+            mask2 = (
+                mask.repeat(1, self.upsample_times, 1)
+                .transpose(-1, -2)
+                .reshape(alphas2.shape[0], -1)
+            )
+            mask2 = mask2.unsqueeze(-1)
+            alphas2 = alphas2 * mask2
+        alphas2 = alphas2.squeeze(-1)
+        token_num2 = alphas2.sum(-1)
+
+        output = output.transpose(1, 2)
+
+        output = self.cif_output(output)
+        alphas = torch.sigmoid(output)
+        alphas = torch.nn.functional.relu(alphas * self.smooth_factor - self.noise_threshold)
+        if mask is not None:
+            mask = mask.transpose(-1, -2).float()
+            alphas = alphas * mask
+        if mask_chunk_predictor is not None:
+            alphas = alphas * mask_chunk_predictor
+        alphas = alphas.squeeze(-1)
+        mask = mask.squeeze(-1)
+        if target_label_length is not None:
+            target_length = target_label_length
+        elif target_label is not None:
+            target_length = (target_label != ignore_id).float().sum(-1)
+        else:
+            target_length = None
+        token_num = alphas.sum(-1)
+
+        if target_length is not None:
+            alphas *= (target_length / token_num)[:, None].repeat(1, alphas.size(1))
+        elif self.tail_threshold > 0.0:
+            hidden, alphas, token_num = self.tail_process_fn(hidden, alphas, token_num, mask=mask)
+
+        acoustic_embeds, cif_peak = cif(hidden, alphas, self.threshold)
+        if target_length is None and self.tail_threshold > 0.0:
+            token_num_int = torch.max(token_num).type(torch.int32).item()
+            acoustic_embeds = acoustic_embeds[:, :token_num_int, :]
+        return acoustic_embeds, token_num, alphas, cif_peak, token_num2
+
+    def get_upsample_timestamp(self, hidden, mask=None, token_num=None):
+        """Get upsample timestamp.
+        
+            Args:
+                hidden: TODO.
+                mask: TODO.
+                token_num: TODO.
+            """
+        h = hidden
+        b = hidden.shape[0]
+        context = h.transpose(1, 2)
+        queries = self.pad(context)
+        output = torch.relu(self.cif_conv1d(queries))
+
+        # alphas2 is an extra head for timestamp prediction
+        if not self.use_cif1_cnn:
+            _output = context
+        else:
+            _output = output
+        if self.upsample_type == "cnn":
+            output2 = self.upsample_cnn(_output)
+            output2 = output2.transpose(1, 2)
+        elif self.upsample_type == "cnn_blstm":
+            output2 = self.upsample_cnn(_output)
+            output2 = output2.transpose(1, 2)
+            output2, (_, _) = self.blstm(output2)
+        elif self.upsample_type == "cnn_attn":
+            output2 = self.upsample_cnn(_output)
+            output2 = output2.transpose(1, 2)
+            output2, _ = self.self_attn(output2, mask)
+        alphas2 = torch.sigmoid(self.cif_output2(output2))
+        alphas2 = torch.nn.functional.relu(alphas2 * self.smooth_factor2 - self.noise_threshold2)
+        # repeat the mask in T demension to match the upsampled length
+        if mask is not None:
+            mask2 = (
+                mask.repeat(1, self.upsample_times, 1)
+                .transpose(-1, -2)
+                .reshape(alphas2.shape[0], -1)
+            )
+            mask2 = mask2.unsqueeze(-1)
+            alphas2 = alphas2 * mask2
+        alphas2 = alphas2.squeeze(-1)
+        _token_num = alphas2.sum(-1)
+        if token_num is not None:
+            alphas2 *= (token_num / _token_num)[:, None].repeat(1, alphas2.size(1))
+        # re-downsample
+        ds_alphas = alphas2.reshape(b, -1, self.upsample_times).sum(-1)
+        ds_cif_peak = cif_wo_hidden(ds_alphas, self.threshold - 1e-4)
+        # upsampled alphas and cif_peak
+        us_alphas = alphas2
+        us_cif_peak = cif_wo_hidden(us_alphas, self.threshold - 1e-4)
+        return ds_alphas, ds_cif_peak, us_alphas, us_cif_peak
+
+    def tail_process_fn(self, hidden, alphas, token_num=None, mask=None):
+        """Tail process fn.
+        
+            Args:
+                hidden: TODO.
+                alphas: TODO.
+                token_num: TODO.
+                mask: TODO.
+            """
+        b, t, d = hidden.size()
+        tail_threshold = self.tail_threshold
+        if mask is not None:
+            zeros_t = torch.zeros((b, 1), dtype=torch.float32, device=alphas.device)
+            ones_t = torch.ones_like(zeros_t)
+            mask_1 = torch.cat([mask, zeros_t], dim=1)
+            mask_2 = torch.cat([ones_t, mask], dim=1)
+            mask = mask_2 - mask_1
+            tail_threshold = mask * tail_threshold
+            alphas = torch.cat([alphas, zeros_t], dim=1)
+            alphas = torch.add(alphas, tail_threshold)
+        else:
+            tail_threshold = torch.tensor([tail_threshold], dtype=alphas.dtype).to(alphas.device)
+            tail_threshold = torch.reshape(tail_threshold, (1, 1))
+            alphas = torch.cat([alphas, tail_threshold], dim=1)
+        zeros = torch.zeros((b, 1, d), dtype=hidden.dtype).to(hidden.device)
+        hidden = torch.cat([hidden, zeros], dim=1)
+        token_num = alphas.sum(dim=-1)
+        token_num_floor = torch.floor(token_num)
+
+        return hidden, alphas, token_num_floor
+
+    def gen_frame_alignments(
+        self, alphas: torch.Tensor = None, encoder_sequence_length: torch.Tensor = None
+    ):
+        """Gen frame alignments.
+        
+            Args:
+                alphas: TODO.
+                encoder_sequence_length: TODO.
+            """
+        batch_size, maximum_length = alphas.size()
+        int_type = torch.int32
+
+        is_training = self.training
+        if is_training:
+            token_num = torch.round(torch.sum(alphas, dim=1)).type(int_type)
+        else:
+            token_num = torch.floor(torch.sum(alphas, dim=1)).type(int_type)
+
+        max_token_num = torch.max(token_num).item()
+
+        alphas_cumsum = torch.cumsum(alphas, dim=1)
+        alphas_cumsum = torch.floor(alphas_cumsum).type(int_type)
+        alphas_cumsum = alphas_cumsum[:, None, :].repeat(1, max_token_num, 1)
+
+        index = torch.ones([batch_size, max_token_num], dtype=int_type)
+        index = torch.cumsum(index, dim=1)
+        index = index[:, :, None].repeat(1, 1, maximum_length).to(alphas_cumsum.device)
+
+        index_div = torch.floor(torch.true_divide(alphas_cumsum, index)).type(int_type)
+        index_div_bool_zeros = index_div.eq(0)
+        index_div_bool_zeros_count = torch.sum(index_div_bool_zeros, dim=-1) + 1
+        index_div_bool_zeros_count = torch.clamp(
+            index_div_bool_zeros_count, 0, encoder_sequence_length.max()
+        )
+        token_num_mask = (~make_pad_mask(token_num, maxlen=max_token_num)).to(token_num.device)
+        index_div_bool_zeros_count *= token_num_mask
+
+        index_div_bool_zeros_count_tile = index_div_bool_zeros_count[:, :, None].repeat(
+            1, 1, maximum_length
+        )
+        ones = torch.ones_like(index_div_bool_zeros_count_tile)
+        zeros = torch.zeros_like(index_div_bool_zeros_count_tile)
+        ones = torch.cumsum(ones, dim=2)
+        cond = index_div_bool_zeros_count_tile == ones
+        index_div_bool_zeros_count_tile = torch.where(cond, zeros, ones)
+
+        index_div_bool_zeros_count_tile_bool = index_div_bool_zeros_count_tile.type(torch.bool)
+        index_div_bool_zeros_count_tile = 1 - index_div_bool_zeros_count_tile_bool.type(int_type)
+        index_div_bool_zeros_count_tile_out = torch.sum(index_div_bool_zeros_count_tile, dim=1)
+        index_div_bool_zeros_count_tile_out = index_div_bool_zeros_count_tile_out.type(int_type)
+        predictor_mask = (
+            (~make_pad_mask(encoder_sequence_length, maxlen=encoder_sequence_length.max()))
+            .type(int_type)
+            .to(encoder_sequence_length.device)
+        )
+        index_div_bool_zeros_count_tile_out = index_div_bool_zeros_count_tile_out * predictor_mask
+
+        predictor_alignments = index_div_bool_zeros_count_tile_out
+        predictor_alignments_length = predictor_alignments.sum(-1).type(
+            encoder_sequence_length.dtype
+        )
+        return predictor_alignments.detach(), predictor_alignments_length.detach()
+
+
+@tables.register("predictor_classes", "CifPredictorV3Export")
+class CifPredictorV3Export(torch.nn.Module):
+    def __init__(self, model, **kwargs):
+        """Initialize CifPredictorV3Export.
+        
+            Args:
+                model: Model instance or model name.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+
+        self.pad = model.pad
+        self.cif_conv1d = model.cif_conv1d
+        self.cif_output = model.cif_output
+        self.threshold = model.threshold
+        self.smooth_factor = model.smooth_factor
+        self.noise_threshold = model.noise_threshold
+        self.tail_threshold = model.tail_threshold
+
+        self.upsample_times = model.upsample_times
+        self.upsample_cnn = model.upsample_cnn
+        self.blstm = model.blstm
+        self.cif_output2 = model.cif_output2
+        self.smooth_factor2 = model.smooth_factor2
+        self.noise_threshold2 = model.noise_threshold2
+
+    def forward(
+        self,
+        hidden: torch.Tensor,
+        mask: torch.Tensor,
+    ):
+        """Forward pass for training.
+        
+            Args:
+                hidden: TODO.
+                mask: TODO.
+            """
+        h = hidden
+        context = h.transpose(1, 2)
+        queries = self.pad(context)
+        output = torch.relu(self.cif_conv1d(queries))
+        output = output.transpose(1, 2)
+
+        output = self.cif_output(output)
+        alphas = torch.sigmoid(output)
+        alphas = torch.nn.functional.relu(alphas * self.smooth_factor - self.noise_threshold)
+        mask = mask.transpose(-1, -2).float()
+        alphas = alphas * mask
+        alphas = alphas.squeeze(-1)
+        token_num = alphas.sum(-1)
+
+        mask = mask.squeeze(-1)
+        hidden, alphas, token_num = self.tail_process_fn(hidden, alphas, mask=mask)
+        acoustic_embeds, cif_peak = cif_export(hidden, alphas, self.threshold)
+
+        return acoustic_embeds, token_num, alphas, cif_peak
+
+    def get_upsample_timestmap(self, hidden, mask=None, token_num=None):
+        """Get upsample timestmap.
+        
+            Args:
+                hidden: TODO.
+                mask: TODO.
+                token_num: TODO.
+            """
+        h = hidden
+        b = hidden.shape[0]
+        context = h.transpose(1, 2)
+
+        # generate alphas2
+        _output = context
+        output2 = self.upsample_cnn(_output)
+        output2 = output2.transpose(1, 2)
+        output2, (_, _) = self.blstm(output2)
+        alphas2 = torch.sigmoid(self.cif_output2(output2))
+        alphas2 = torch.nn.functional.relu(alphas2 * self.smooth_factor2 - self.noise_threshold2)
+
+        mask = (
+            mask.repeat(1, self.upsample_times, 1).transpose(-1, -2).reshape(alphas2.shape[0], -1)
+        )
+        mask = mask.unsqueeze(-1)
+        alphas2 = alphas2 * mask
+        alphas2 = alphas2.squeeze(-1)
+        _token_num = alphas2.sum(-1)
+        alphas2 *= (token_num / _token_num)[:, None].repeat(1, alphas2.size(1))
+        # upsampled alphas and cif_peak
+        us_alphas = alphas2
+        us_cif_peak = cif_wo_hidden_export(us_alphas, self.threshold - 1e-4)
+        return us_alphas, us_cif_peak
+
+    def tail_process_fn(self, hidden, alphas, token_num=None, mask=None):
+        """Tail process fn.
+        
+            Args:
+                hidden: TODO.
+                alphas: TODO.
+                token_num: TODO.
+                mask: TODO.
+            """
+        b, t, d = hidden.size()
+        tail_threshold = self.tail_threshold
+
+        zeros_t = torch.zeros((b, 1), dtype=torch.float32, device=alphas.device)
+        ones_t = torch.ones_like(zeros_t)
+
+        mask_1 = torch.cat([mask, zeros_t], dim=1)
+        mask_2 = torch.cat([ones_t, mask], dim=1)
+        mask = mask_2 - mask_1
+        tail_threshold = mask * tail_threshold
+        alphas = torch.cat([alphas, zeros_t], dim=1)
+        alphas = torch.add(alphas, tail_threshold)
+
+        zeros = torch.zeros((b, 1, d), dtype=hidden.dtype).to(hidden.device)
+        hidden = torch.cat([hidden, zeros], dim=1)
+        token_num = alphas.sum(dim=-1)
+        token_num_floor = torch.floor(token_num)
+
+        return hidden, alphas, token_num_floor
+
+
+@torch.jit.script
+def cif_export(hidden, alphas, threshold: float):
+    """Cif export.
+    
+        Args:
+            hidden: TODO.
+            alphas: TODO.
+            threshold: TODO.
+        """
+    batch_size, len_time, hidden_size = hidden.size()
+    threshold = torch.tensor([threshold], dtype=alphas.dtype).to(alphas.device)
+
+    # loop varss
+    integrate = torch.zeros([batch_size], dtype=alphas.dtype, device=hidden.device)
+    frame = torch.zeros([batch_size, hidden_size], dtype=hidden.dtype, device=hidden.device)
+    # intermediate vars along time
+    list_fires = []
+    list_frames = []
+
+    for t in range(len_time):
+        alpha = alphas[:, t]
+        distribution_completion = (
+            torch.ones([batch_size], dtype=alphas.dtype, device=hidden.device) - integrate
+        )
+
+        integrate += alpha
+        list_fires.append(integrate)
+
+        fire_place = integrate >= threshold
+        integrate = torch.where(
+            fire_place,
+            integrate - torch.ones([batch_size], dtype=alphas.dtype, device=hidden.device),
+            integrate,
+        )
+        cur = torch.where(fire_place, distribution_completion, alpha)
+        remainds = alpha - cur
+
+        frame += cur[:, None] * hidden[:, t, :]
+        list_frames.append(frame)
+        frame = torch.where(
+            fire_place[:, None].repeat(1, hidden_size), remainds[:, None] * hidden[:, t, :], frame
+        )
+
+    fires = torch.stack(list_fires, 1)
+    frames = torch.stack(list_frames, 1)
+
+    fire_idxs = fires >= threshold
+    frame_fires = torch.zeros_like(hidden)
+    max_label_len = frames[0, fire_idxs[0]].size(0)
+    for b in range(batch_size):
+        frame_fire = frames[b, fire_idxs[b]]
+        frame_len = frame_fire.size(0)
+        frame_fires[b, :frame_len, :] = frame_fire
+
+        if frame_len >= max_label_len:
+            max_label_len = frame_len
+    frame_fires = frame_fires[:, :max_label_len, :]
+    return frame_fires, fires
+
+
+@torch.jit.script
+def cif_wo_hidden_export(alphas, threshold: float):
+    """Cif wo hidden export.
+    
+        Args:
+            alphas: TODO.
+            threshold: TODO.
+        """
+    batch_size, len_time = alphas.size()
+
+    # loop varss
+    integrate = torch.zeros([batch_size], dtype=alphas.dtype, device=alphas.device)
+    # intermediate vars along time
+    list_fires = []
+
+    for t in range(len_time):
+        alpha = alphas[:, t]
+
+        integrate += alpha
+        list_fires.append(integrate)
+
+        fire_place = integrate >= threshold
+        integrate = torch.where(
+            fire_place,
+            integrate - torch.ones([batch_size], device=alphas.device) * threshold,
+            integrate,
+        )
+
+    fires = torch.stack(list_fires, 1)
+    return fires
+

funasr/models/bicif_paraformer/export_meta.py

@@ -0,0 +1,104 @@
+#!/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 torch
+import types
+
+from funasr.register import tables
+
+
+def export_rebuild_model(model, **kwargs):
+    """Export rebuild model.
+    
+        Args:
+            model: Model instance or model name.
+            **kwargs: Additional keyword arguments.
+        """
+    is_onnx = kwargs.get("type", "onnx") == "onnx"
+    encoder_class = tables.encoder_classes.get(kwargs["encoder"] + "Export")
+    model.encoder = encoder_class(model.encoder, onnx=is_onnx)
+
+    predictor_class = tables.predictor_classes.get(kwargs["predictor"] + "Export")
+    model.predictor = predictor_class(model.predictor, onnx=is_onnx)
+
+    decoder_class = tables.decoder_classes.get(kwargs["decoder"] + "Export")
+    model.decoder = decoder_class(model.decoder, onnx=is_onnx)
+
+    from funasr.utils.torch_function import sequence_mask
+
+    model.make_pad_mask = sequence_mask(kwargs["max_seq_len"], flip=False)
+
+    model.forward = types.MethodType(export_forward, model)
+    model.export_dummy_inputs = types.MethodType(export_dummy_inputs, model)
+    model.export_input_names = types.MethodType(export_input_names, model)
+    model.export_output_names = types.MethodType(export_output_names, model)
+    model.export_dynamic_axes = types.MethodType(export_dynamic_axes, model)
+    
+    model.export_name = "model"
+
+    return model
+
+
+def export_forward(
+    self,
+    speech: torch.Tensor,
+    speech_lengths: torch.Tensor,
+):
+    # a. To device
+    """Export forward.
+    
+        Args:
+            speech: Speech audio tensor, shape (batch, time).
+            speech_lengths: Length of each speech sample.
+        """
+    batch = {"speech": speech, "speech_lengths": speech_lengths}
+
+    enc, enc_len = self.encoder(**batch)
+    mask = self.make_pad_mask(enc_len)[:, None, :]
+    pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = self.predictor(enc, mask)
+    pre_token_length = pre_token_length.round().type(torch.int32)
+
+    decoder_out, _ = self.decoder(enc, enc_len, pre_acoustic_embeds, pre_token_length)
+    decoder_out = torch.log_softmax(decoder_out, dim=-1)
+
+    # get predicted timestamps
+    us_alphas, us_cif_peak = self.predictor.get_upsample_timestmap(enc, mask, pre_token_length)
+
+    return decoder_out, pre_token_length, us_alphas, us_cif_peak
+
+
+def export_dummy_inputs(self):
+    """Export dummy inputs."""
+    speech = torch.randn(2, 30, 560)
+    speech_lengths = torch.tensor([6, 30], dtype=torch.int32)
+    return (speech, speech_lengths)
+
+
+def export_input_names(self):
+    """Export input names."""
+    return ["speech", "speech_lengths"]
+
+
+def export_output_names(self):
+    """Export output names."""
+    return ["logits", "token_num", "us_alphas", "us_cif_peak"]
+
+
+def export_dynamic_axes(self):
+    """Export dynamic axes."""
+    return {
+        "speech": {0: "batch_size", 1: "feats_length"},
+        "speech_lengths": {
+            0: "batch_size",
+        },
+        "logits": {0: "batch_size", 1: "logits_length"},
+        "us_alphas": {0: "batch_size", 1: "alphas_length"},
+        "us_cif_peak": {0: "batch_size", 1: "alphas_length"},
+    }
+
+
+def export_name(self):
+    """Export name."""
+    return "model.onnx"

funasr/models/bicif_paraformer/model.py

@@ -0,0 +1,441 @@
+#!/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 copy
+import time
+import torch
+import logging
+from contextlib import contextmanager
+from distutils.version import LooseVersion
+from typing import Dict, List, Optional, Tuple
+
+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.train_utils.device_funcs import force_gatherable
+from funasr.models.transformer.utils.add_sos_eos import add_sos_eos
+from funasr.utils.timestamp_tools import ts_prediction_lfr6_standard
+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.train_utils.device_funcs import to_device
+
+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):
+        """Autocast.
+        
+            Args:
+                enabled: TODO.
+            """
+        yield
+
+
+@tables.register("model_classes", "BiCifParaformer")
+class BiCifParaformer(Paraformer):
+    """BiCifParaformer: Paraformer with Bidirectional CIF for Timestamp Prediction.
+
+    Extends Paraformer with a second CIF predictor that provides accurate
+    character-level timestamp prediction alongside ASR. Uses bidirectional
+    information flow for better alignment between audio frames and text tokens.
+
+    Reference:
+        - FunASR: A Fundamental End-to-End Speech Recognition Toolkit (https://arxiv.org/abs/2305.11013)
+        - Achieving timestamp prediction while recognizing with non-autoregressive end-to-end ASR model
+          (https://arxiv.org/abs/2301.12343)
+
+    Output:
+        {"key": str, "text": str, "timestamp": [[start_ms, end_ms], ...]}
+
+    Author: Speech Lab of DAMO Academy, Alibaba Group
+    """
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        """Initialize BiCifParaformer.
+        
+            Args:
+                *args: Variable positional arguments.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__(*args, **kwargs)
+
+    def _calc_pre2_loss(
+        self,
+        encoder_out: torch.Tensor,
+        encoder_out_lens: torch.Tensor,
+        ys_pad: torch.Tensor,
+        ys_pad_lens: torch.Tensor,
+    ):
+        """Internal: calc pre2 loss.
+        
+            Args:
+                encoder_out: Encoder output tensor.
+                encoder_out_lens: Encoder output lengths.
+                ys_pad: TODO.
+                ys_pad_lens: Lengths of ys_pad.
+            """
+        encoder_out_mask = (
+            ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]
+        ).to(encoder_out.device)
+        if self.predictor_bias == 1:
+            _, ys_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id)
+            ys_pad_lens = ys_pad_lens + self.predictor_bias
+        _, _, _, _, pre_token_length2 = self.predictor(
+            encoder_out, ys_pad, encoder_out_mask, ignore_id=self.ignore_id
+        )
+
+        # loss_pre = self.criterion_pre(ys_pad_lens.type_as(pre_token_length), pre_token_length)
+        loss_pre2 = self.criterion_pre(ys_pad_lens.type_as(pre_token_length2), pre_token_length2)
+
+        return loss_pre2
+
+    def _calc_att_loss(
+        self,
+        encoder_out: torch.Tensor,
+        encoder_out_lens: torch.Tensor,
+        ys_pad: torch.Tensor,
+        ys_pad_lens: torch.Tensor,
+    ):
+        """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.
+            """
+        encoder_out_mask = (
+            ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]
+        ).to(encoder_out.device)
+        if self.predictor_bias == 1:
+            _, ys_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id)
+            ys_pad_lens = ys_pad_lens + self.predictor_bias
+        pre_acoustic_embeds, pre_token_length, _, pre_peak_index, _ = self.predictor(
+            encoder_out, ys_pad, encoder_out_mask, ignore_id=self.ignore_id
+        )
+
+        # 0. sampler
+        decoder_out_1st = None
+        if self.sampling_ratio > 0.0:
+            sematic_embeds, decoder_out_1st = self.sampler(
+                encoder_out, encoder_out_lens, ys_pad, ys_pad_lens, pre_acoustic_embeds
+            )
+        else:
+            sematic_embeds = pre_acoustic_embeds
+
+        # 1. Forward decoder
+        decoder_outs = self.decoder(encoder_out, encoder_out_lens, sematic_embeds, 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,
+        )
+        loss_pre = self.criterion_pre(ys_pad_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_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, loss_pre
+
+    def calc_predictor(self, encoder_out, encoder_out_lens):
+        """Calc predictor.
+        
+            Args:
+                encoder_out: Encoder output tensor.
+                encoder_out_lens: Encoder output lengths.
+            """
+        encoder_out_mask = (
+            ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]
+        ).to(encoder_out.device)
+        pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index, pre_token_length2 = (
+            self.predictor(encoder_out, None, encoder_out_mask, ignore_id=self.ignore_id)
+        )
+        return pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index
+
+    def calc_predictor_timestamp(self, encoder_out, encoder_out_lens, token_num):
+        """Calc predictor timestamp.
+        
+            Args:
+                encoder_out: Encoder output tensor.
+                encoder_out_lens: Encoder output lengths.
+                token_num: TODO.
+            """
+        encoder_out_mask = (
+            ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]
+        ).to(encoder_out.device)
+        ds_alphas, ds_cif_peak, us_alphas, us_peaks = self.predictor.get_upsample_timestamp(
+            encoder_out, encoder_out_mask, token_num
+        )
+        return ds_alphas, ds_cif_peak, us_alphas, us_peaks
+
+    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]:
+        """Frontend + 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, loss_pre = self._calc_att_loss(
+            encoder_out, encoder_out_lens, text, text_lengths
+        )
+
+        loss_pre2 = self._calc_pre2_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
+                + loss_pre2 * self.predictor_weight * 0.5
+            )
+        else:
+            loss = (
+                self.ctc_weight * loss_ctc
+                + (1 - self.ctc_weight) * loss_att
+                + loss_pre * self.predictor_weight
+                + loss_pre2 * self.predictor_weight * 0.5
+            )
+
+        # 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_pre2"] = loss_pre2.detach().cpu()
+
+        stats["loss"] = torch.clone(loss.detach())
+
+        # force_gatherable: to-device and to-tensor if scalar for DataParallel
+        if self.length_normalized_loss:
+            batch_size = int((text_lengths + self.predictor_bias).sum())
+
+        loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device)
+        return loss, stats, weight
+
+    def inference(
+        self,
+        data_in,
+        data_lengths=None,
+        key: list = None,
+        tokenizer=None,
+        frontend=None,
+        **kwargs,
+    ):
+
+        # init beamsearch
+        """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.
+            """
+        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 and (is_use_lm or is_use_ctc):
+            logging.info("enable beam_search")
+            self.init_beam_search(**kwargs)
+            self.nbest = kwargs.get("nbest", 1)
+
+        meta_data = {}
+        # if isinstance(data_in, torch.Tensor):  # fbank
+        #     speech, speech_lengths = data_in, data_lengths
+        #     if len(speech.shape) < 3:
+        #         speech = speech[None, :, :]
+        #     if speech_lengths is None:
+        #         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)
+        )
+        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
+        encoder_out, encoder_out_lens = self.encode(speech, speech_lengths)
+        if isinstance(encoder_out, tuple):
+            encoder_out = encoder_out[0]
+
+        # predictor
+        predictor_outs = self.calc_predictor(encoder_out, encoder_out_lens)
+        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 []
+        decoder_outs = self.cal_decoder_with_predictor(
+            encoder_out, encoder_out_lens, pre_acoustic_embeds, pre_token_length
+        )
+        decoder_out, ys_pad_lens = decoder_outs[0], decoder_outs[1]
+
+        # BiCifParaformer, test no bias cif2
+        _, _, us_alphas, us_peaks = self.calc_predictor_timestamp(
+            encoder_out, encoder_out_lens, pre_token_length
+        )
+
+        results = []
+        b, n, d = decoder_out.size()
+        for i in range(b):
+            x = encoder_out[i, : encoder_out_lens[i], :]
+            am_scores = decoder_out[i, : pre_token_length[i], :]
+            if self.beam_search is not None:
+                nbest_hyps = self.beam_search(
+                    x=x,
+                    am_scores=am_scores,
+                    maxlenratio=kwargs.get("maxlenratio", 0.0),
+                    minlenratio=kwargs.get("minlenratio", 0.0),
+                )
+
+                nbest_hyps = nbest_hyps[: self.nbest]
+            else:
+
+                yseq = am_scores.argmax(dim=-1)
+                score = am_scores.max(dim=-1)[0]
+                score = torch.sum(score, dim=-1)
+                # pad with mask tokens to ensure compatibility with sos/eos tokens
+                yseq = torch.tensor([self.sos] + yseq.tolist() + [self.eos], device=yseq.device)
+                nbest_hyps = [Hypothesis(yseq=yseq, score=score)]
+            for nbest_idx, hyp in enumerate(nbest_hyps):
+                ibest_writer = None
+                if kwargs.get("output_dir") is not None:
+                    if not hasattr(self, "writer"):
+                        self.writer = DatadirWriter(kwargs.get("output_dir"))
+                    ibest_writer = self.writer[f"{nbest_idx+1}best_recog"]
+
+                # 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
+                    )
+                )
+
+                if tokenizer is not None:
+                    # Change integer-ids to tokens
+                    token = tokenizer.ids2tokens(token_int)
+                    text = tokenizer.tokens2text(token)
+
+                    _, timestamp = ts_prediction_lfr6_standard(
+                        us_alphas[i][: encoder_out_lens[i] * 3],
+                        us_peaks[i][: encoder_out_lens[i] * 3],
+                        copy.copy(token),
+                        vad_offset=kwargs.get("begin_time", 0),
+                    )
+
+                    text_postprocessed, time_stamp_postprocessed, word_lists = (
+                        postprocess_utils.sentence_postprocess(token, timestamp)
+                    )
+
+                    result_i = {
+                        "key": key[i],
+                        "text": text_postprocessed,
+                        "timestamp": time_stamp_postprocessed,
+                    }
+
+                    if ibest_writer is not None:
+                        ibest_writer["token"][key[i]] = " ".join(token)
+                        # ibest_writer["text"][key[i]] = text
+                        ibest_writer["timestamp"][key[i]] = time_stamp_postprocessed
+                        ibest_writer["text"][key[i]] = text_postprocessed
+                else:
+                    result_i = {"key": key[i], "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

funasr/models/bicif_paraformer/template.yaml

@@ -0,0 +1,134 @@
+# 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: funasr.models.paraformer.model:Paraformer
+model: BiCifParaformer
+model_conf:
+    ctc_weight: 0.0
+    lsm_weight: 0.1
+    length_normalized_loss: true
+    predictor_weight: 1.0
+    predictor_bias: 1
+    sampling_ratio: 0.75
+
+# encoder
+encoder: SANMEncoder
+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: ParaformerSANMDecoder
+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: CifPredictorV3
+predictor_conf:
+    idim: 512
+    threshold: 1.0
+    l_order: 1
+    r_order: 1
+    tail_threshold: 0.45
+    smooth_factor2: 0.25
+    noise_threshold2: 0.01
+    upsample_times: 3
+    use_cif1_cnn: false
+    upsample_type: cnn_blstm
+
+# 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
+  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

funasr/models/branchformer/cgmlp.py

@@ -0,0 +1,150 @@
+"""MLP with convolutional gating (cgMLP) definition.
+
+References:
+    https://openreview.net/forum?id=RA-zVvZLYIy
+    https://arxiv.org/abs/2105.08050
+
+"""
+
+import torch
+
+from funasr.models.transformer.utils.nets_utils import get_activation
+from funasr.models.transformer.layer_norm import LayerNorm
+
+
+class ConvolutionalSpatialGatingUnit(torch.nn.Module):
+    """Convolutional Spatial Gating Unit (CSGU)."""
+
+    def __init__(
+        self,
+        size: int,
+        kernel_size: int,
+        dropout_rate: float,
+        use_linear_after_conv: bool,
+        gate_activation: str,
+    ):
+        """Initialize ConvolutionalSpatialGatingUnit.
+        
+            Args:
+                size: TODO.
+                kernel_size: Size/dimension parameter.
+                dropout_rate: TODO.
+                use_linear_after_conv: TODO.
+                gate_activation: TODO.
+            """
+        super().__init__()
+
+        n_channels = size // 2  # split input channels
+        self.norm = LayerNorm(n_channels)
+        self.conv = torch.nn.Conv1d(
+            n_channels,
+            n_channels,
+            kernel_size,
+            1,
+            (kernel_size - 1) // 2,
+            groups=n_channels,
+        )
+        if use_linear_after_conv:
+            self.linear = torch.nn.Linear(n_channels, n_channels)
+        else:
+            self.linear = None
+
+        if gate_activation == "identity":
+            self.act = torch.nn.Identity()
+        else:
+            self.act = get_activation(gate_activation)
+
+        self.dropout = torch.nn.Dropout(dropout_rate)
+
+    def espnet_initialization_fn(self):
+        """Espnet initialization fn."""
+        torch.nn.init.normal_(self.conv.weight, std=1e-6)
+        torch.nn.init.ones_(self.conv.bias)
+        if self.linear is not None:
+            torch.nn.init.normal_(self.linear.weight, std=1e-6)
+            torch.nn.init.ones_(self.linear.bias)
+
+    def forward(self, x, gate_add=None):
+        """Forward method
+
+        Args:
+            x (torch.Tensor): (N, T, D)
+            gate_add (torch.Tensor): (N, T, D/2)
+
+        Returns:
+            out (torch.Tensor): (N, T, D/2)
+        """
+
+        x_r, x_g = x.chunk(2, dim=-1)
+
+        x_g = self.norm(x_g)  # (N, T, D/2)
+        x_g = self.conv(x_g.transpose(1, 2)).transpose(1, 2)  # (N, T, D/2)
+        if self.linear is not None:
+            x_g = self.linear(x_g)
+
+        if gate_add is not None:
+            x_g = x_g + gate_add
+
+        x_g = self.act(x_g)
+        out = x_r * x_g  # (N, T, D/2)
+        out = self.dropout(out)
+        return out
+
+
+class ConvolutionalGatingMLP(torch.nn.Module):
+    """Convolutional Gating MLP (cgMLP)."""
+
+    def __init__(
+        self,
+        size: int,
+        linear_units: int,
+        kernel_size: int,
+        dropout_rate: float,
+        use_linear_after_conv: bool,
+        gate_activation: str,
+    ):
+        """Initialize ConvolutionalGatingMLP.
+        
+            Args:
+                size: TODO.
+                linear_units: TODO.
+                kernel_size: Size/dimension parameter.
+                dropout_rate: TODO.
+                use_linear_after_conv: TODO.
+                gate_activation: TODO.
+            """
+        super().__init__()
+
+        self.channel_proj1 = torch.nn.Sequential(
+            torch.nn.Linear(size, linear_units), torch.nn.GELU()
+        )
+        self.csgu = ConvolutionalSpatialGatingUnit(
+            size=linear_units,
+            kernel_size=kernel_size,
+            dropout_rate=dropout_rate,
+            use_linear_after_conv=use_linear_after_conv,
+            gate_activation=gate_activation,
+        )
+        self.channel_proj2 = torch.nn.Linear(linear_units // 2, size)
+
+    def forward(self, x, mask):
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+                mask: TODO.
+            """
+        if isinstance(x, tuple):
+            xs_pad, pos_emb = x
+        else:
+            xs_pad, pos_emb = x, None
+
+        xs_pad = self.channel_proj1(xs_pad)  # size -> linear_units
+        xs_pad = self.csgu(xs_pad)  # linear_units -> linear_units/2
+        xs_pad = self.channel_proj2(xs_pad)  # linear_units/2 -> size
+
+        if pos_emb is not None:
+            out = (xs_pad, pos_emb)
+        else:
+            out = xs_pad
+        return out

funasr/models/branchformer/encoder.py

@@ -0,0 +1,564 @@
+# Copyright 2022 Yifan Peng (Carnegie Mellon University)
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+"""Branchformer encoder definition.
+
+Reference:
+    Yifan Peng, Siddharth Dalmia, Ian Lane, and Shinji Watanabe,
+    “Branchformer: Parallel MLP-Attention Architectures to Capture
+    Local and Global Context for Speech Recognition and Understanding,”
+    in Proceedings of ICML, 2022.
+
+"""
+
+import logging
+from typing import List, Optional, Tuple, Union
+
+import numpy
+import torch
+import torch.nn as nn
+
+from funasr.models.branchformer.cgmlp import ConvolutionalGatingMLP
+from funasr.models.branchformer.fastformer import FastSelfAttention
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+from funasr.models.transformer.attention import (  # noqa: H301
+    LegacyRelPositionMultiHeadedAttention,
+    MultiHeadedAttention,
+    RelPositionMultiHeadedAttention,
+)
+from funasr.models.transformer.embedding import (  # noqa: H301
+    LegacyRelPositionalEncoding,
+    PositionalEncoding,
+    RelPositionalEncoding,
+    ScaledPositionalEncoding,
+)
+from funasr.models.transformer.layer_norm import LayerNorm
+from funasr.models.transformer.utils.repeat import repeat
+from funasr.models.transformer.utils.subsampling import (
+    Conv2dSubsampling,
+    Conv2dSubsampling2,
+    Conv2dSubsampling6,
+    Conv2dSubsampling8,
+    TooShortUttError,
+    check_short_utt,
+)
+
+from funasr.register import tables
+
+
+class BranchformerEncoderLayer(torch.nn.Module):
+    """Branchformer encoder layer module.
+
+    Args:
+        size (int): model dimension
+        attn: standard self-attention or efficient attention, optional
+        cgmlp: ConvolutionalGatingMLP, optional
+        dropout_rate (float): dropout probability
+        merge_method (str): concat, learned_ave, fixed_ave
+        cgmlp_weight (float): weight of the cgmlp branch, between 0 and 1,
+            used if merge_method is fixed_ave
+        attn_branch_drop_rate (float): probability of dropping the attn branch,
+            used if merge_method is learned_ave
+        stochastic_depth_rate (float): stochastic depth probability
+    """
+
+    def __init__(
+        self,
+        size: int,
+        attn: Optional[torch.nn.Module],
+        cgmlp: Optional[torch.nn.Module],
+        dropout_rate: float,
+        merge_method: str,
+        cgmlp_weight: float = 0.5,
+        attn_branch_drop_rate: float = 0.0,
+        stochastic_depth_rate: float = 0.0,
+    ):
+        """Initialize BranchformerEncoderLayer.
+        
+            Args:
+                size: TODO.
+                attn: TODO.
+                cgmlp: TODO.
+                dropout_rate: TODO.
+                merge_method: TODO.
+                cgmlp_weight: TODO.
+                attn_branch_drop_rate: TODO.
+                stochastic_depth_rate: TODO.
+            """
+        super().__init__()
+        assert (attn is not None) or (cgmlp is not None), "At least one branch should be valid"
+
+        self.size = size
+        self.attn = attn
+        self.cgmlp = cgmlp
+        self.merge_method = merge_method
+        self.cgmlp_weight = cgmlp_weight
+        self.attn_branch_drop_rate = attn_branch_drop_rate
+        self.stochastic_depth_rate = stochastic_depth_rate
+        self.use_two_branches = (attn is not None) and (cgmlp is not None)
+
+        if attn is not None:
+            self.norm_mha = LayerNorm(size)  # for the MHA module
+        if cgmlp is not None:
+            self.norm_mlp = LayerNorm(size)  # for the MLP module
+        self.norm_final = LayerNorm(size)  # for the final output of the block
+
+        self.dropout = torch.nn.Dropout(dropout_rate)
+
+        if self.use_two_branches:
+            if merge_method == "concat":
+                self.merge_proj = torch.nn.Linear(size + size, size)
+
+            elif merge_method == "learned_ave":
+                # attention-based pooling for two branches
+                self.pooling_proj1 = torch.nn.Linear(size, 1)
+                self.pooling_proj2 = torch.nn.Linear(size, 1)
+
+                # linear projections for calculating merging weights
+                self.weight_proj1 = torch.nn.Linear(size, 1)
+                self.weight_proj2 = torch.nn.Linear(size, 1)
+
+                # linear projection after weighted average
+                self.merge_proj = torch.nn.Linear(size, size)
+
+            elif merge_method == "fixed_ave":
+                assert 0.0 <= cgmlp_weight <= 1.0, "cgmlp weight should be between 0.0 and 1.0"
+
+                # remove the other branch if only one branch is used
+                if cgmlp_weight == 0.0:
+                    self.use_two_branches = False
+                    self.cgmlp = None
+                    self.norm_mlp = None
+                elif cgmlp_weight == 1.0:
+                    self.use_two_branches = False
+                    self.attn = None
+                    self.norm_mha = None
+
+                # linear projection after weighted average
+                self.merge_proj = torch.nn.Linear(size, size)
+
+            else:
+                raise ValueError(f"unknown merge method: {merge_method}")
+
+        else:
+            self.merge_proj = torch.nn.Identity()
+
+    def forward(self, x_input, mask, cache=None):
+        """Compute encoded features.
+
+        Args:
+            x_input (Union[Tuple, torch.Tensor]): Input tensor w/ or w/o pos emb.
+                - w/ pos emb: Tuple of tensors [(#batch, time, size), (1, time, size)].
+                - w/o pos emb: Tensor (#batch, time, size).
+            mask (torch.Tensor): Mask tensor for the input (#batch, 1, 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).
+        """
+
+        if cache is not None:
+            raise NotImplementedError("cache is not None, which is not tested")
+
+        if isinstance(x_input, tuple):
+            x, pos_emb = x_input[0], x_input[1]
+        else:
+            x, pos_emb = x_input, None
+
+        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)
+            if pos_emb is not None:
+                return (x, pos_emb), mask
+            return x, mask
+
+        # Two branches
+        x1 = x
+        x2 = x
+
+        # Branch 1: multi-headed attention module
+        if self.attn is not None:
+            x1 = self.norm_mha(x1)
+
+            if isinstance(self.attn, FastSelfAttention):
+                x_att = self.attn(x1, mask)
+            else:
+                if pos_emb is not None:
+                    x_att = self.attn(x1, x1, x1, pos_emb, mask)
+                else:
+                    x_att = self.attn(x1, x1, x1, mask)
+
+            x1 = self.dropout(x_att)
+
+        # Branch 2: convolutional gating mlp
+        if self.cgmlp is not None:
+            x2 = self.norm_mlp(x2)
+
+            if pos_emb is not None:
+                x2 = (x2, pos_emb)
+            x2 = self.cgmlp(x2, mask)
+            if isinstance(x2, tuple):
+                x2 = x2[0]
+
+            x2 = self.dropout(x2)
+
+        # Merge two branches
+        if self.use_two_branches:
+            if self.merge_method == "concat":
+                x = x + stoch_layer_coeff * self.dropout(
+                    self.merge_proj(torch.cat([x1, x2], dim=-1))
+                )
+            elif self.merge_method == "learned_ave":
+                if (
+                    self.training
+                    and self.attn_branch_drop_rate > 0
+                    and torch.rand(1).item() < self.attn_branch_drop_rate
+                ):
+                    # Drop the attn branch
+                    w1, w2 = 0.0, 1.0
+                else:
+                    # branch1
+                    score1 = (
+                        self.pooling_proj1(x1).transpose(1, 2) / self.size**0.5
+                    )  # (batch, 1, time)
+                    if mask is not None:
+                        min_value = float(
+                            numpy.finfo(torch.tensor(0, dtype=score1.dtype).numpy().dtype).min
+                        )
+                        score1 = score1.masked_fill(mask.eq(0), min_value)
+                        score1 = torch.softmax(score1, dim=-1).masked_fill(mask.eq(0), 0.0)
+                    else:
+                        score1 = torch.softmax(score1, dim=-1)
+                    pooled1 = torch.matmul(score1, x1).squeeze(1)  # (batch, size)
+                    weight1 = self.weight_proj1(pooled1)  # (batch, 1)
+
+                    # branch2
+                    score2 = (
+                        self.pooling_proj2(x2).transpose(1, 2) / self.size**0.5
+                    )  # (batch, 1, time)
+                    if mask is not None:
+                        min_value = float(
+                            numpy.finfo(torch.tensor(0, dtype=score2.dtype).numpy().dtype).min
+                        )
+                        score2 = score2.masked_fill(mask.eq(0), min_value)
+                        score2 = torch.softmax(score2, dim=-1).masked_fill(mask.eq(0), 0.0)
+                    else:
+                        score2 = torch.softmax(score2, dim=-1)
+                    pooled2 = torch.matmul(score2, x2).squeeze(1)  # (batch, size)
+                    weight2 = self.weight_proj2(pooled2)  # (batch, 1)
+
+                    # normalize weights of two branches
+                    merge_weights = torch.softmax(
+                        torch.cat([weight1, weight2], dim=-1), dim=-1
+                    )  # (batch, 2)
+                    merge_weights = merge_weights.unsqueeze(-1).unsqueeze(-1)  # (batch, 2, 1, 1)
+                    w1, w2 = merge_weights[:, 0], merge_weights[:, 1]  # (batch, 1, 1)
+
+                x = x + stoch_layer_coeff * self.dropout(self.merge_proj(w1 * x1 + w2 * x2))
+            elif self.merge_method == "fixed_ave":
+                x = x + stoch_layer_coeff * self.dropout(
+                    self.merge_proj((1.0 - self.cgmlp_weight) * x1 + self.cgmlp_weight * x2)
+                )
+            else:
+                raise RuntimeError(f"unknown merge method: {self.merge_method}")
+        else:
+            if self.attn is None:
+                x = x + stoch_layer_coeff * self.dropout(self.merge_proj(x2))
+            elif self.cgmlp is None:
+                x = x + stoch_layer_coeff * self.dropout(self.merge_proj(x1))
+            else:
+                # This should not happen
+                raise RuntimeError("Both branches are not None, which is unexpected.")
+
+        x = self.norm_final(x)
+
+        if pos_emb is not None:
+            return (x, pos_emb), mask
+
+        return x, mask
+
+
+@tables.register("encoder_classes", "BranchformerEncoder")
+class BranchformerEncoder(nn.Module):
+    """Branchformer encoder module."""
+
+    def __init__(
+        self,
+        input_size: int,
+        output_size: int = 256,
+        use_attn: bool = True,
+        attention_heads: int = 4,
+        attention_layer_type: str = "rel_selfattn",
+        pos_enc_layer_type: str = "rel_pos",
+        rel_pos_type: str = "latest",
+        use_cgmlp: bool = True,
+        cgmlp_linear_units: int = 2048,
+        cgmlp_conv_kernel: int = 31,
+        use_linear_after_conv: bool = False,
+        gate_activation: str = "identity",
+        merge_method: str = "concat",
+        cgmlp_weight: Union[float, List[float]] = 0.5,
+        attn_branch_drop_rate: Union[float, List[float]] = 0.0,
+        num_blocks: int = 12,
+        dropout_rate: float = 0.1,
+        positional_dropout_rate: float = 0.1,
+        attention_dropout_rate: float = 0.0,
+        input_layer: Optional[str] = "conv2d",
+        zero_triu: bool = False,
+        padding_idx: int = -1,
+        stochastic_depth_rate: Union[float, List[float]] = 0.0,
+    ):
+        """Initialize BranchformerEncoder.
+        
+            Args:
+                input_size: Size/dimension parameter.
+                output_size: Size/dimension parameter.
+                use_attn: TODO.
+                attention_heads: TODO.
+                attention_layer_type: TODO.
+                pos_enc_layer_type: TODO.
+                rel_pos_type: TODO.
+                use_cgmlp: TODO.
+                cgmlp_linear_units: TODO.
+                cgmlp_conv_kernel: TODO.
+                use_linear_after_conv: TODO.
+                gate_activation: TODO.
+                merge_method: TODO.
+                cgmlp_weight: TODO.
+                attn_branch_drop_rate: TODO.
+                num_blocks: TODO.
+                dropout_rate: TODO.
+                positional_dropout_rate: TODO.
+                attention_dropout_rate: TODO.
+                input_layer: TODO.
+                zero_triu: TODO.
+                padding_idx: TODO.
+                stochastic_depth_rate: TODO.
+            """
+        super().__init__()
+        self._output_size = output_size
+
+        if rel_pos_type == "legacy":
+            if pos_enc_layer_type == "rel_pos":
+                pos_enc_layer_type = "legacy_rel_pos"
+            if attention_layer_type == "rel_selfattn":
+                attention_layer_type = "legacy_rel_selfattn"
+        elif rel_pos_type == "latest":
+            assert attention_layer_type != "legacy_rel_selfattn"
+            assert pos_enc_layer_type != "legacy_rel_pos"
+        else:
+            raise ValueError("unknown rel_pos_type: " + rel_pos_type)
+
+        if pos_enc_layer_type == "abs_pos":
+            pos_enc_class = PositionalEncoding
+        elif pos_enc_layer_type == "scaled_abs_pos":
+            pos_enc_class = ScaledPositionalEncoding
+        elif pos_enc_layer_type == "rel_pos":
+            assert attention_layer_type == "rel_selfattn"
+            pos_enc_class = RelPositionalEncoding
+        elif pos_enc_layer_type == "legacy_rel_pos":
+            assert attention_layer_type == "legacy_rel_selfattn"
+            pos_enc_class = LegacyRelPositionalEncoding
+            logging.warning("Using legacy_rel_pos and it will be deprecated in the future.")
+        else:
+            raise ValueError("unknown pos_enc_layer: " + pos_enc_layer_type)
+
+        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),
+                pos_enc_class(output_size, positional_dropout_rate),
+            )
+        elif input_layer == "conv2d":
+            self.embed = Conv2dSubsampling(
+                input_size,
+                output_size,
+                dropout_rate,
+                pos_enc_class(output_size, positional_dropout_rate),
+            )
+        elif input_layer == "conv2d2":
+            self.embed = Conv2dSubsampling2(
+                input_size,
+                output_size,
+                dropout_rate,
+                pos_enc_class(output_size, positional_dropout_rate),
+            )
+        elif input_layer == "conv2d6":
+            self.embed = Conv2dSubsampling6(
+                input_size,
+                output_size,
+                dropout_rate,
+                pos_enc_class(output_size, positional_dropout_rate),
+            )
+        elif input_layer == "conv2d8":
+            self.embed = Conv2dSubsampling8(
+                input_size,
+                output_size,
+                dropout_rate,
+                pos_enc_class(output_size, positional_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 isinstance(input_layer, torch.nn.Module):
+            self.embed = torch.nn.Sequential(
+                input_layer,
+                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)
+        else:
+            raise ValueError("unknown input_layer: " + input_layer)
+
+        if attention_layer_type == "selfattn":
+            encoder_selfattn_layer = MultiHeadedAttention
+            encoder_selfattn_layer_args = (
+                attention_heads,
+                output_size,
+                attention_dropout_rate,
+            )
+        elif attention_layer_type == "legacy_rel_selfattn":
+            assert pos_enc_layer_type == "legacy_rel_pos"
+            encoder_selfattn_layer = LegacyRelPositionMultiHeadedAttention
+            encoder_selfattn_layer_args = (
+                attention_heads,
+                output_size,
+                attention_dropout_rate,
+            )
+            logging.warning("Using legacy_rel_selfattn and it will be deprecated in the future.")
+        elif attention_layer_type == "rel_selfattn":
+            assert pos_enc_layer_type == "rel_pos"
+            encoder_selfattn_layer = RelPositionMultiHeadedAttention
+            encoder_selfattn_layer_args = (
+                attention_heads,
+                output_size,
+                attention_dropout_rate,
+                zero_triu,
+            )
+        elif attention_layer_type == "fast_selfattn":
+            assert pos_enc_layer_type in ["abs_pos", "scaled_abs_pos"]
+            encoder_selfattn_layer = FastSelfAttention
+            encoder_selfattn_layer_args = (
+                output_size,
+                attention_heads,
+                attention_dropout_rate,
+            )
+        else:
+            raise ValueError("unknown encoder_attn_layer: " + attention_layer_type)
+
+        cgmlp_layer = ConvolutionalGatingMLP
+        cgmlp_layer_args = (
+            output_size,
+            cgmlp_linear_units,
+            cgmlp_conv_kernel,
+            dropout_rate,
+            use_linear_after_conv,
+            gate_activation,
+        )
+
+        if isinstance(stochastic_depth_rate, float):
+            stochastic_depth_rate = [stochastic_depth_rate] * num_blocks
+        if len(stochastic_depth_rate) != num_blocks:
+            raise ValueError(
+                f"Length of stochastic_depth_rate ({len(stochastic_depth_rate)}) "
+                f"should be equal to num_blocks ({num_blocks})"
+            )
+
+        if isinstance(cgmlp_weight, float):
+            cgmlp_weight = [cgmlp_weight] * num_blocks
+        if len(cgmlp_weight) != num_blocks:
+            raise ValueError(
+                f"Length of cgmlp_weight ({len(cgmlp_weight)}) should be equal to "
+                f"num_blocks ({num_blocks})"
+            )
+
+        if isinstance(attn_branch_drop_rate, float):
+            attn_branch_drop_rate = [attn_branch_drop_rate] * num_blocks
+        if len(attn_branch_drop_rate) != num_blocks:
+            raise ValueError(
+                f"Length of attn_branch_drop_rate ({len(attn_branch_drop_rate)}) "
+                f"should be equal to num_blocks ({num_blocks})"
+            )
+
+        self.encoders = repeat(
+            num_blocks,
+            lambda lnum: BranchformerEncoderLayer(
+                output_size,
+                encoder_selfattn_layer(*encoder_selfattn_layer_args) if use_attn else None,
+                cgmlp_layer(*cgmlp_layer_args) if use_cgmlp else None,
+                dropout_rate,
+                merge_method,
+                cgmlp_weight[lnum],
+                attn_branch_drop_rate[lnum],
+                stochastic_depth_rate[lnum],
+            ),
+        )
+        self.after_norm = LayerNorm(output_size)
+
+    def output_size(self) -> int:
+        """Output size."""
+        return self._output_size
+
+    def forward(
+        self,
+        xs_pad: torch.Tensor,
+        ilens: torch.Tensor,
+        prev_states: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        """Calculate forward propagation.
+
+        Args:
+            xs_pad (torch.Tensor): Input tensor (#batch, L, input_size).
+            ilens (torch.Tensor): Input length (#batch).
+            prev_states (torch.Tensor): Not to be used now.
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, L, output_size).
+            torch.Tensor: Output length (#batch).
+            torch.Tensor: Not to be used now.
+
+        """
+
+        masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device)
+
+        if (
+            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)
+        elif self.embed is not None:
+            xs_pad = self.embed(xs_pad)
+
+        xs_pad, masks = self.encoders(xs_pad, masks)
+
+        if isinstance(xs_pad, tuple):
+            xs_pad = xs_pad[0]
+
+        xs_pad = self.after_norm(xs_pad)
+        olens = masks.squeeze(1).sum(1)
+        return xs_pad, olens, None

funasr/models/branchformer/fastformer.py

@@ -0,0 +1,152 @@
+"""Fastformer attention definition.
+
+Reference:
+    Wu et al., "Fastformer: Additive Attention Can Be All You Need"
+    https://arxiv.org/abs/2108.09084
+    https://github.com/wuch15/Fastformer
+
+"""
+
+import numpy
+import torch
+
+
+class FastSelfAttention(torch.nn.Module):
+    """Fast self-attention used in Fastformer."""
+
+    def __init__(
+        self,
+        size,
+        attention_heads,
+        dropout_rate,
+    ):
+        """Initialize FastSelfAttention.
+        
+            Args:
+                size: TODO.
+                attention_heads: TODO.
+                dropout_rate: TODO.
+            """
+        super().__init__()
+        if size % attention_heads != 0:
+            raise ValueError(
+                f"Hidden size ({size}) is not an integer multiple "
+                f"of attention heads ({attention_heads})"
+            )
+        self.attention_head_size = size // attention_heads
+        self.num_attention_heads = attention_heads
+
+        self.query = torch.nn.Linear(size, size)
+        self.query_att = torch.nn.Linear(size, attention_heads)
+        self.key = torch.nn.Linear(size, size)
+        self.key_att = torch.nn.Linear(size, attention_heads)
+        self.transform = torch.nn.Linear(size, size)
+        self.dropout = torch.nn.Dropout(dropout_rate)
+
+    def espnet_initialization_fn(self):
+        """Espnet initialization fn."""
+        self.apply(self.init_weights)
+
+    def init_weights(self, module):
+        """Init weights.
+        
+            Args:
+                module: TODO.
+            """
+        if isinstance(module, torch.nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+        if isinstance(module, torch.nn.Linear) and module.bias is not None:
+            module.bias.data.zero_()
+
+    def transpose_for_scores(self, x):
+        """Reshape and transpose to compute scores.
+
+        Args:
+            x: (batch, time, size = n_heads * attn_dim)
+
+        Returns:
+            (batch, n_heads, time, attn_dim)
+        """
+
+        new_x_shape = x.shape[:-1] + (
+            self.num_attention_heads,
+            self.attention_head_size,
+        )
+        return x.reshape(*new_x_shape).transpose(1, 2)
+
+    def forward(self, xs_pad, mask):
+        """Forward method.
+
+        Args:
+            xs_pad: (batch, time, size = n_heads * attn_dim)
+            mask: (batch, 1, time), nonpadding is 1, padding is 0
+
+        Returns:
+            torch.Tensor: (batch, time, size)
+        """
+
+        batch_size, seq_len, _ = xs_pad.shape
+        mixed_query_layer = self.query(xs_pad)  # (batch, time, size)
+        mixed_key_layer = self.key(xs_pad)  # (batch, time, size)
+
+        if mask is not None:
+            mask = mask.eq(0)  # padding is 1, nonpadding is 0
+
+        # (batch, n_heads, time)
+        query_for_score = (
+            self.query_att(mixed_query_layer).transpose(1, 2) / self.attention_head_size**0.5
+        )
+        if mask is not None:
+            min_value = float(
+                numpy.finfo(torch.tensor(0, dtype=query_for_score.dtype).numpy().dtype).min
+            )
+            query_for_score = query_for_score.masked_fill(mask, min_value)
+            query_weight = torch.softmax(query_for_score, dim=-1).masked_fill(mask, 0.0)
+        else:
+            query_weight = torch.softmax(query_for_score, dim=-1)
+
+        query_weight = query_weight.unsqueeze(2)  # (batch, n_heads, 1, time)
+        query_layer = self.transpose_for_scores(
+            mixed_query_layer
+        )  # (batch, n_heads, time, attn_dim)
+
+        pooled_query = (
+            torch.matmul(query_weight, query_layer)
+            .transpose(1, 2)
+            .reshape(-1, 1, self.num_attention_heads * self.attention_head_size)
+        )  # (batch, 1, size = n_heads * attn_dim)
+        pooled_query = self.dropout(pooled_query)
+        pooled_query_repeat = pooled_query.repeat(1, seq_len, 1)  # (batch, time, size)
+
+        mixed_query_key_layer = mixed_key_layer * pooled_query_repeat  # (batch, time, size)
+
+        # (batch, n_heads, time)
+        query_key_score = (
+            self.key_att(mixed_query_key_layer) / self.attention_head_size**0.5
+        ).transpose(1, 2)
+        if mask is not None:
+            min_value = float(
+                numpy.finfo(torch.tensor(0, dtype=query_key_score.dtype).numpy().dtype).min
+            )
+            query_key_score = query_key_score.masked_fill(mask, min_value)
+            query_key_weight = torch.softmax(query_key_score, dim=-1).masked_fill(mask, 0.0)
+        else:
+            query_key_weight = torch.softmax(query_key_score, dim=-1)
+
+        query_key_weight = query_key_weight.unsqueeze(2)  # (batch, n_heads, 1, time)
+        key_layer = self.transpose_for_scores(
+            mixed_query_key_layer
+        )  # (batch, n_heads, time, attn_dim)
+        pooled_key = torch.matmul(query_key_weight, key_layer)  # (batch, n_heads, 1, attn_dim)
+        pooled_key = self.dropout(pooled_key)
+
+        # NOTE: value = query, due to param sharing
+        weighted_value = (pooled_key * query_layer).transpose(
+            1, 2
+        )  # (batch, time, n_heads, attn_dim)
+        weighted_value = weighted_value.reshape(
+            weighted_value.shape[:-2] + (self.num_attention_heads * self.attention_head_size,)
+        )  # (batch, time, size)
+        weighted_value = self.dropout(self.transform(weighted_value)) + mixed_query_layer
+
+        return weighted_value

funasr/models/branchformer/model.py

@@ -0,0 +1,29 @@
+import logging
+
+from funasr.models.transformer.model import Transformer
+from funasr.register import tables
+
+
+@tables.register("model_classes", "Branchformer")
+class Branchformer(Transformer):
+    """Branchformer: Parallel branch encoder architecture.
+
+    Uses parallel branches of self-attention and convolution that are
+    merged via concatenation. Alternative to Conformer with similar accuracy.
+
+    Inherits Transformer pipeline for training and inference.
+    """
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+
+        """Initialize Branchformer.
+        
+            Args:
+                *args: Variable positional arguments.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__(*args, **kwargs)

funasr/models/branchformer/template.yaml

@@ -0,0 +1,116 @@
+# 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: Branchformer
+model_conf:
+    ctc_weight: 0.3
+    lsm_weight: 0.1     # label smoothing option
+    length_normalized_loss: false
+
+# encoder
+encoder: BranchformerEncoder
+encoder_conf:
+    output_size: 256
+    use_attn: true
+    attention_heads: 4
+    attention_layer_type: rel_selfattn
+    pos_enc_layer_type: rel_pos
+    rel_pos_type: latest
+    use_cgmlp: true
+    cgmlp_linear_units: 2048
+    cgmlp_conv_kernel: 31
+    use_linear_after_conv: false
+    gate_activation: identity
+    merge_method: concat
+    cgmlp_weight: 0.5               # used only if merge_method is "fixed_ave"
+    attn_branch_drop_rate: 0.0      # used only if merge_method is "learned_ave"
+    num_blocks: 24
+    dropout_rate: 0.1
+    positional_dropout_rate: 0.1
+    attention_dropout_rate: 0.1
+    input_layer: conv2d
+    stochastic_depth_rate: 0.0
+
+# decoder
+decoder: TransformerDecoder
+decoder_conf:
+    attention_heads: 4
+    linear_units: 2048
+    num_blocks: 6
+    dropout_rate: 0.1
+    positional_dropout_rate: 0.1
+    self_attention_dropout_rate: 0.
+    src_attention_dropout_rate: 0.
+
+
+# frontend related
+frontend: WavFrontend
+frontend_conf:
+    fs: 16000
+    window: hamming
+    n_mels: 80
+    frame_length: 25
+    frame_shift: 10
+    dither: 0.0
+    lfr_m: 1
+    lfr_n: 1
+
+specaug: SpecAug
+specaug_conf:
+    apply_time_warp: true
+    time_warp_window: 5
+    time_warp_mode: bicubic
+    apply_freq_mask: true
+    freq_mask_width_range:
+    - 0
+    - 30
+    num_freq_mask: 2
+    apply_time_mask: true
+    time_mask_width_range:
+    - 0
+    - 40
+    num_time_mask: 2
+
+train_conf:
+  accum_grad: 1
+  grad_clip: 5
+  max_epoch: 150
+  keep_nbest_models: 10
+  log_interval: 50
+
+optim: adam
+optim_conf:
+   lr: 0.001
+   weight_decay: 0.000001
+scheduler: warmuplr
+scheduler_conf:
+   warmup_steps: 35000
+
+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: 4
+
+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

funasr/models/campplus/cluster_backend.py

@@ -0,0 +1,268 @@
+#!/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)
+# Modified from 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker)
+
+import scipy
+import torch
+import sklearn
+import numpy as np
+
+from sklearn.cluster._kmeans import k_means
+from sklearn.cluster import HDBSCAN
+
+
+class SpectralCluster:
+    r"""A spectral clustering mehtod using unnormalized Laplacian of affinity matrix.
+    This implementation is adapted from https://github.com/speechbrain/speechbrain.
+    """
+
+    def __init__(self, min_num_spks=1, max_num_spks=15, pval=0.022):
+        """Initialize SpectralCluster.
+        
+            Args:
+                min_num_spks: TODO.
+                max_num_spks: TODO.
+                pval: TODO.
+            """
+        self.min_num_spks = min_num_spks
+        self.max_num_spks = max_num_spks
+        self.pval = pval
+
+    def __call__(self, X, oracle_num=None):
+        # Similarity matrix computation
+        """Internal: call  .
+        
+            Args:
+                X: TODO.
+                oracle_num: TODO.
+            """
+        sim_mat = self.get_sim_mat(X)
+
+        # Refining similarity matrix with pval
+        prunned_sim_mat = self.p_pruning(sim_mat)
+
+        # Symmetrization
+        sym_prund_sim_mat = 0.5 * (prunned_sim_mat + prunned_sim_mat.T)
+
+        # Laplacian calculation
+        laplacian = self.get_laplacian(sym_prund_sim_mat)
+
+        # Get Spectral Embeddings
+        emb, num_of_spk = self.get_spec_embs(laplacian, oracle_num)
+
+        # Perform clustering
+        labels = self.cluster_embs(emb, num_of_spk)
+
+        return labels
+
+    def get_sim_mat(self, X):
+        # Cosine similarities
+        """Get sim mat.
+        
+            Args:
+                X: TODO.
+            """
+        M = sklearn.metrics.pairwise.cosine_similarity(X, X)
+        return M
+
+    def p_pruning(self, A):
+        """P pruning.
+        
+            Args:
+                A: TODO.
+            """
+        if A.shape[0] * self.pval < 6:
+            pval = 6.0 / A.shape[0]
+        else:
+            pval = self.pval
+
+        n_elems = int((1 - pval) * A.shape[0])
+
+        # For each row in a affinity matrix
+        for i in range(A.shape[0]):
+            low_indexes = np.argsort(A[i, :])
+            low_indexes = low_indexes[0:n_elems]
+
+            # Replace smaller similarity values by 0s
+            A[i, low_indexes] = 0
+        return A
+
+    def get_laplacian(self, M):
+        """Get laplacian.
+        
+            Args:
+                M: TODO.
+            """
+        M[np.diag_indices(M.shape[0])] = 0
+        D = np.sum(np.abs(M), axis=1)
+        D = np.diag(D)
+        L = D - M
+        return L
+
+    def get_spec_embs(self, L, k_oracle=None):
+        """Get spec embs.
+        
+            Args:
+                L: TODO.
+                k_oracle: TODO.
+            """
+        lambdas, eig_vecs = scipy.linalg.eigh(L)
+
+        if k_oracle is not None:
+            num_of_spk = k_oracle
+        else:
+            lambda_gap_list = self.getEigenGaps(
+                lambdas[self.min_num_spks - 1 : self.max_num_spks + 1]
+            )
+            num_of_spk = np.argmax(lambda_gap_list) + self.min_num_spks
+
+        emb = eig_vecs[:, :num_of_spk]
+        return emb, num_of_spk
+
+    def cluster_embs(self, emb, k):
+        """Cluster embs.
+        
+            Args:
+                emb: TODO.
+                k: TODO.
+            """
+        _, labels, _ = k_means(emb, k)
+        return labels
+
+    def getEigenGaps(self, eig_vals):
+        """Geteigengaps.
+        
+            Args:
+                eig_vals: TODO.
+            """
+        eig_vals_gap_list = []
+        for i in range(len(eig_vals) - 1):
+            gap = float(eig_vals[i + 1]) - float(eig_vals[i])
+            eig_vals_gap_list.append(gap)
+        return eig_vals_gap_list
+
+
+class UmapHdbscan:
+    r"""
+    Reference:
+    - Siqi Zheng, Hongbin Suo. Reformulating Speaker Diarization as Community Detection With
+      Emphasis On Topological Structure. ICASSP2022
+    """
+
+    def __init__(
+        self, n_neighbors=20, n_components=60, min_samples=10, min_cluster_size=10, metric="cosine"
+    ):
+        """Initialize UmapHdbscan.
+        
+            Args:
+                n_neighbors: TODO.
+                n_components: TODO.
+                min_samples: TODO.
+                min_cluster_size: Size/dimension parameter.
+                metric: TODO.
+            """
+        self.n_neighbors = n_neighbors
+        self.n_components = n_components
+        self.min_samples = min_samples
+        self.min_cluster_size = min_cluster_size
+        self.metric = metric
+
+    def __call__(self, X):
+        """Internal: call  .
+        
+            Args:
+                X: TODO.
+            """
+        import umap.umap_ as umap
+
+        umap_X = umap.UMAP(
+            n_neighbors=self.n_neighbors,
+            min_dist=0.0,
+            n_components=min(self.n_components, X.shape[0] - 2),
+            metric=self.metric,
+        ).fit_transform(X)
+        labels = HDBSCAN(
+            min_samples=self.min_samples,
+            min_cluster_size=self.min_cluster_size,
+            allow_single_cluster=True,
+        ).fit_predict(umap_X)
+        return labels
+
+
+class ClusterBackend(torch.nn.Module):
+    r"""Perfom clustering for input embeddings and output the labels.
+    Args:
+        model_dir: A model dir.
+        model_config: The model config.
+    """
+
+    def __init__(self, merge_thr=0.78):
+        """Initialize ClusterBackend.
+        
+            Args:
+                merge_thr: TODO.
+            """
+        super().__init__()
+        self.model_config = {"merge_thr": merge_thr}
+        # self.other_config = kwargs
+
+        self.spectral_cluster = SpectralCluster()
+        self.umap_hdbscan_cluster = UmapHdbscan()
+
+    def forward(self, X, **params):
+        # clustering and return the labels
+        """Forward pass for training.
+        
+            Args:
+                X: TODO.
+                **params: Additional keyword arguments.
+            """
+        k = params["oracle_num"] if "oracle_num" in params else None
+        assert len(X.shape) == 2, "modelscope error: the shape of input should be [N, C]"
+        if X.shape[0] < 20:
+            return np.zeros(X.shape[0], dtype="int")
+        if X.shape[0] < 2048 or k is not None:
+            # unexpected corner case
+            labels = self.spectral_cluster(X, k)
+        else:
+            labels = self.umap_hdbscan_cluster(X)
+
+        if k is None and "merge_thr" in self.model_config:
+            labels = self.merge_by_cos(labels, X, self.model_config["merge_thr"])
+
+        return labels
+
+    def merge_by_cos(self, labels, embs, cos_thr):
+        # merge the similar speakers by cosine similarity
+        """Merge by cos.
+        
+            Args:
+                labels: TODO.
+                embs: TODO.
+                cos_thr: TODO.
+            """
+        assert cos_thr > 0 and cos_thr <= 1
+        while True:
+            spk_num = labels.max() + 1
+            if spk_num == 1:
+                break
+            spk_center = []
+            for i in range(spk_num):
+                spk_emb = embs[labels == i].mean(0)
+                spk_center.append(spk_emb)
+            assert len(spk_center) > 0
+            spk_center = np.stack(spk_center, axis=0)
+            norm_spk_center = spk_center / np.linalg.norm(spk_center, axis=1, keepdims=True)
+            affinity = np.matmul(norm_spk_center, norm_spk_center.T)
+            affinity = np.triu(affinity, 1)
+            spks = np.unravel_index(np.argmax(affinity), affinity.shape)
+            if affinity[spks] < cos_thr:
+                break
+            for i in range(len(labels)):
+                if labels[i] == spks[1]:
+                    labels[i] = spks[0]
+                elif labels[i] > spks[1]:
+                    labels[i] -= 1
+        return labels

funasr/models/campplus/components.py

@@ -0,0 +1,450 @@
+#!/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)
+# Modified from 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker)
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint as cp
+
+
+class BasicResBlock(torch.nn.Module):
+    expansion = 1
+
+    def __init__(self, in_planes, planes, stride=1):
+        """Initialize BasicResBlock.
+        
+            Args:
+                in_planes: TODO.
+                planes: TODO.
+                stride: TODO.
+            """
+        super(BasicResBlock, self).__init__()
+        self.conv1 = torch.nn.Conv2d(
+            in_planes, planes, kernel_size=3, stride=(stride, 1), padding=1, bias=False
+        )
+        self.bn1 = torch.nn.BatchNorm2d(planes)
+        self.conv2 = torch.nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = torch.nn.BatchNorm2d(planes)
+
+        self.shortcut = torch.nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            self.shortcut = torch.nn.Sequential(
+                torch.nn.Conv2d(
+                    in_planes,
+                    self.expansion * planes,
+                    kernel_size=1,
+                    stride=(stride, 1),
+                    bias=False,
+                ),
+                torch.nn.BatchNorm2d(self.expansion * planes),
+            )
+
+    def forward(self, x):
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+            """
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+        return out
+
+
+class FCM(torch.nn.Module):
+    def __init__(self, block=BasicResBlock, num_blocks=[2, 2], m_channels=32, feat_dim=80):
+        """Initialize FCM.
+        
+            Args:
+                block: TODO.
+                num_blocks: TODO.
+                m_channels: TODO.
+                feat_dim: Size/dimension parameter.
+            """
+        super(FCM, self).__init__()
+        self.in_planes = m_channels
+        self.conv1 = torch.nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = torch.nn.BatchNorm2d(m_channels)
+
+        self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=2)
+        self.layer2 = self._make_layer(block, m_channels, num_blocks[0], stride=2)
+
+        self.conv2 = torch.nn.Conv2d(
+            m_channels, m_channels, kernel_size=3, stride=(2, 1), padding=1, bias=False
+        )
+        self.bn2 = torch.nn.BatchNorm2d(m_channels)
+        self.out_channels = m_channels * (feat_dim // 8)
+
+    def _make_layer(self, block, planes, num_blocks, stride):
+        """Internal: make layer.
+        
+            Args:
+                block: TODO.
+                planes: TODO.
+                num_blocks: TODO.
+                stride: TODO.
+            """
+        strides = [stride] + [1] * (num_blocks - 1)
+        layers = []
+        for stride in strides:
+            layers.append(block(self.in_planes, planes, stride))
+            self.in_planes = planes * block.expansion
+        return torch.nn.Sequential(*layers)
+
+    def forward(self, x):
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+            """
+        x = x.unsqueeze(1)
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.layer1(out)
+        out = self.layer2(out)
+        out = F.relu(self.bn2(self.conv2(out)))
+
+        shape = out.shape
+        out = out.reshape(shape[0], shape[1] * shape[2], shape[3])
+        return out
+
+
+def get_nonlinear(config_str, channels):
+    """Get nonlinear.
+    
+        Args:
+            config_str: TODO.
+            channels: TODO.
+        """
+    nonlinear = torch.nn.Sequential()
+    for name in config_str.split("-"):
+        if name == "relu":
+            nonlinear.add_module("relu", torch.nn.ReLU(inplace=True))
+        elif name == "prelu":
+            nonlinear.add_module("prelu", torch.nn.PReLU(channels))
+        elif name == "batchnorm":
+            nonlinear.add_module("batchnorm", torch.nn.BatchNorm1d(channels))
+        elif name == "batchnorm_":
+            nonlinear.add_module("batchnorm", torch.nn.BatchNorm1d(channels, affine=False))
+        else:
+            raise ValueError("Unexpected module ({}).".format(name))
+    return nonlinear
+
+
+def statistics_pooling(x, dim=-1, keepdim=False, unbiased=True, eps=1e-2):
+    """Statistics pooling.
+    
+        Args:
+            x: TODO.
+            dim: TODO.
+            keepdim: TODO.
+            unbiased: TODO.
+            eps: TODO.
+        """
+    mean = x.mean(dim=dim)
+    std = x.std(dim=dim, unbiased=unbiased)
+    stats = torch.cat([mean, std], dim=-1)
+    if keepdim:
+        stats = stats.unsqueeze(dim=dim)
+    return stats
+
+
+class StatsPool(torch.nn.Module):
+    def forward(self, x):
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+            """
+        return statistics_pooling(x)
+
+
+class TDNNLayer(torch.nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        bias=False,
+        config_str="batchnorm-relu",
+    ):
+        """Initialize TDNNLayer.
+        
+            Args:
+                in_channels: TODO.
+                out_channels: TODO.
+                kernel_size: Size/dimension parameter.
+                stride: TODO.
+                padding: TODO.
+                dilation: TODO.
+                bias: TODO.
+                config_str: TODO.
+            """
+        super(TDNNLayer, self).__init__()
+        if padding < 0:
+            assert (
+                kernel_size % 2 == 1
+            ), "Expect equal paddings, but got even kernel size ({})".format(kernel_size)
+            padding = (kernel_size - 1) // 2 * dilation
+        self.linear = torch.nn.Conv1d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            bias=bias,
+        )
+        self.nonlinear = get_nonlinear(config_str, out_channels)
+
+    def forward(self, x):
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+            """
+        x = self.linear(x)
+        x = self.nonlinear(x)
+        return x
+
+
+class CAMLayer(torch.nn.Module):
+    def __init__(
+        self, bn_channels, out_channels, kernel_size, stride, padding, dilation, bias, reduction=2
+    ):
+        """Initialize CAMLayer.
+        
+            Args:
+                bn_channels: TODO.
+                out_channels: TODO.
+                kernel_size: Size/dimension parameter.
+                stride: TODO.
+                padding: TODO.
+                dilation: TODO.
+                bias: TODO.
+                reduction: TODO.
+            """
+        super(CAMLayer, self).__init__()
+        self.linear_local = torch.nn.Conv1d(
+            bn_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            bias=bias,
+        )
+        self.linear1 = torch.nn.Conv1d(bn_channels, bn_channels // reduction, 1)
+        self.relu = torch.nn.ReLU(inplace=True)
+        self.linear2 = torch.nn.Conv1d(bn_channels // reduction, out_channels, 1)
+        self.sigmoid = torch.nn.Sigmoid()
+
+    def forward(self, x):
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+            """
+        y = self.linear_local(x)
+        context = x.mean(-1, keepdim=True) + self.seg_pooling(x)
+        context = self.relu(self.linear1(context))
+        m = self.sigmoid(self.linear2(context))
+        return y * m
+
+    def seg_pooling(self, x, seg_len=100, stype="avg"):
+        """Seg pooling.
+        
+            Args:
+                x: TODO.
+                seg_len: TODO.
+                stype: TODO.
+            """
+        if stype == "avg":
+            seg = F.avg_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True)
+        elif stype == "max":
+            seg = F.max_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True)
+        else:
+            raise ValueError("Wrong segment pooling type.")
+        shape = seg.shape
+        seg = seg.unsqueeze(-1).expand(*shape, seg_len).reshape(*shape[:-1], -1)
+        seg = seg[..., : x.shape[-1]]
+        return seg
+
+
+class CAMDenseTDNNLayer(torch.nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        bn_channels,
+        kernel_size,
+        stride=1,
+        dilation=1,
+        bias=False,
+        config_str="batchnorm-relu",
+        memory_efficient=False,
+    ):
+        """Initialize CAMDenseTDNNLayer.
+        
+            Args:
+                in_channels: TODO.
+                out_channels: TODO.
+                bn_channels: TODO.
+                kernel_size: Size/dimension parameter.
+                stride: TODO.
+                dilation: TODO.
+                bias: TODO.
+                config_str: TODO.
+                memory_efficient: TODO.
+            """
+        super(CAMDenseTDNNLayer, self).__init__()
+        assert kernel_size % 2 == 1, "Expect equal paddings, but got even kernel size ({})".format(
+            kernel_size
+        )
+        padding = (kernel_size - 1) // 2 * dilation
+        self.memory_efficient = memory_efficient
+        self.nonlinear1 = get_nonlinear(config_str, in_channels)
+        self.linear1 = torch.nn.Conv1d(in_channels, bn_channels, 1, bias=False)
+        self.nonlinear2 = get_nonlinear(config_str, bn_channels)
+        self.cam_layer = CAMLayer(
+            bn_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            bias=bias,
+        )
+
+    def bn_function(self, x):
+        """Bn function.
+        
+            Args:
+                x: TODO.
+            """
+        return self.linear1(self.nonlinear1(x))
+
+    def forward(self, x):
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+            """
+        if self.training and self.memory_efficient:
+            x = cp.checkpoint(self.bn_function, x)
+        else:
+            x = self.bn_function(x)
+        x = self.cam_layer(self.nonlinear2(x))
+        return x
+
+
+class CAMDenseTDNNBlock(torch.nn.ModuleList):
+    def __init__(
+        self,
+        num_layers,
+        in_channels,
+        out_channels,
+        bn_channels,
+        kernel_size,
+        stride=1,
+        dilation=1,
+        bias=False,
+        config_str="batchnorm-relu",
+        memory_efficient=False,
+    ):
+        """Initialize CAMDenseTDNNBlock.
+        
+            Args:
+                num_layers: TODO.
+                in_channels: TODO.
+                out_channels: TODO.
+                bn_channels: TODO.
+                kernel_size: Size/dimension parameter.
+                stride: TODO.
+                dilation: TODO.
+                bias: TODO.
+                config_str: TODO.
+                memory_efficient: TODO.
+            """
+        super(CAMDenseTDNNBlock, self).__init__()
+        for i in range(num_layers):
+            layer = CAMDenseTDNNLayer(
+                in_channels=in_channels + i * out_channels,
+                out_channels=out_channels,
+                bn_channels=bn_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                dilation=dilation,
+                bias=bias,
+                config_str=config_str,
+                memory_efficient=memory_efficient,
+            )
+            self.add_module("tdnnd%d" % (i + 1), layer)
+
+    def forward(self, x):
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+            """
+        for layer in self:
+            x = torch.cat([x, layer(x)], dim=1)
+        return x
+
+
+class TransitLayer(torch.nn.Module):
+    def __init__(self, in_channels, out_channels, bias=True, config_str="batchnorm-relu"):
+        """Initialize TransitLayer.
+        
+            Args:
+                in_channels: TODO.
+                out_channels: TODO.
+                bias: TODO.
+                config_str: TODO.
+            """
+        super(TransitLayer, self).__init__()
+        self.nonlinear = get_nonlinear(config_str, in_channels)
+        self.linear = torch.nn.Conv1d(in_channels, out_channels, 1, bias=bias)
+
+    def forward(self, x):
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+            """
+        x = self.nonlinear(x)
+        x = self.linear(x)
+        return x
+
+
+class DenseLayer(torch.nn.Module):
+    def __init__(self, in_channels, out_channels, bias=False, config_str="batchnorm-relu"):
+        """Initialize DenseLayer.
+        
+            Args:
+                in_channels: TODO.
+                out_channels: TODO.
+                bias: TODO.
+                config_str: TODO.
+            """
+        super(DenseLayer, self).__init__()
+        self.linear = torch.nn.Conv1d(in_channels, out_channels, 1, bias=bias)
+        self.nonlinear = get_nonlinear(config_str, out_channels)
+
+    def forward(self, x):
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+            """
+        if len(x.shape) == 2:
+            x = self.linear(x.unsqueeze(dim=-1)).squeeze(dim=-1)
+        else:
+            x = self.linear(x)
+        x = self.nonlinear(x)
+        return x

funasr/models/campplus/model.py

@@ -0,0 +1,195 @@
+#!/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)
+# Modified from 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker)
+
+import time
+import torch
+import numpy as np
+from collections import OrderedDict
+from contextlib import contextmanager
+from distutils.version import LooseVersion
+
+from funasr.register import tables
+from funasr.models.campplus.utils import extract_feature
+from funasr.utils.load_utils import load_audio_text_image_video
+from funasr.models.campplus.components import (
+    DenseLayer,
+    StatsPool,
+    TDNNLayer,
+    CAMDenseTDNNBlock,
+    TransitLayer,
+    get_nonlinear,
+    FCM,
+)
+
+
+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):
+        """Autocast.
+        
+            Args:
+                enabled: TODO.
+            """
+        yield
+
+
+@tables.register("model_classes", "CAMPPlus")
+class CAMPPlus(torch.nn.Module):
+    """CAM++ Speaker Verification Model.
+
+    Extracts fixed-dimensional speaker embeddings from variable-length audio.
+    Used for speaker verification and speaker diarization pipelines.
+
+    Output: 192-dimensional speaker embedding per utterance.
+    """
+
+    def __init__(
+        self,
+        feat_dim=80,
+        embedding_size=192,
+        growth_rate=32,
+        bn_size=4,
+        init_channels=128,
+        config_str="batchnorm-relu",
+        memory_efficient=True,
+        output_level="segment",
+        **kwargs,
+    ):
+        """Initialize CAMPPlus.
+        
+            Args:
+                feat_dim: Size/dimension parameter.
+                embedding_size: Size/dimension parameter.
+                growth_rate: TODO.
+                bn_size: Size/dimension parameter.
+                init_channels: TODO.
+                config_str: TODO.
+                memory_efficient: TODO.
+                output_level: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+
+        self.head = FCM(feat_dim=feat_dim)
+        channels = self.head.out_channels
+        self.output_level = output_level
+
+        self.xvector = torch.nn.Sequential(
+            OrderedDict(
+                [
+                    (
+                        "tdnn",
+                        TDNNLayer(
+                            channels,
+                            init_channels,
+                            5,
+                            stride=2,
+                            dilation=1,
+                            padding=-1,
+                            config_str=config_str,
+                        ),
+                    ),
+                ]
+            )
+        )
+        channels = init_channels
+        for i, (num_layers, kernel_size, dilation) in enumerate(
+            zip((12, 24, 16), (3, 3, 3), (1, 2, 2))
+        ):
+            block = CAMDenseTDNNBlock(
+                num_layers=num_layers,
+                in_channels=channels,
+                out_channels=growth_rate,
+                bn_channels=bn_size * growth_rate,
+                kernel_size=kernel_size,
+                dilation=dilation,
+                config_str=config_str,
+                memory_efficient=memory_efficient,
+            )
+            self.xvector.add_module("block%d" % (i + 1), block)
+            channels = channels + num_layers * growth_rate
+            self.xvector.add_module(
+                "transit%d" % (i + 1),
+                TransitLayer(channels, channels // 2, bias=False, config_str=config_str),
+            )
+            channels //= 2
+
+        self.xvector.add_module("out_nonlinear", get_nonlinear(config_str, channels))
+
+        if self.output_level == "segment":
+            self.xvector.add_module("stats", StatsPool())
+            self.xvector.add_module(
+                "dense", DenseLayer(channels * 2, embedding_size, config_str="batchnorm_")
+            )
+        else:
+            assert (
+                self.output_level == "frame"
+            ), "`output_level` should be set to 'segment' or 'frame'. "
+
+        for m in self.modules():
+            if isinstance(m, (torch.nn.Conv1d, torch.nn.Linear)):
+                torch.nn.init.kaiming_normal_(m.weight.data)
+                if m.bias is not None:
+                    torch.nn.init.zeros_(m.bias)
+
+    def forward(self, x):
+        """Extract speaker embedding from fbank features.
+
+        Args:
+            x (Tensor): Input fbank features, shape (batch, time, feat_dim).
+
+        Returns:
+            Tensor: Speaker embedding, shape (batch, embedding_size) for segment level,
+                or (batch, time, channels) for frame level.
+        """
+        x = x.permute(0, 2, 1)  # (B,T,F) => (B,F,T)
+        x = self.head(x)
+        x = self.xvector(x)
+        if self.output_level == "frame":
+            x = x.transpose(1, 2)
+        return x
+
+    def inference(
+        self,
+        data_in,
+        data_lengths=None,
+        key: list = None,
+        tokenizer=None,
+        frontend=None,
+        **kwargs,
+    ):
+        """Run speaker embedding extraction on audio input.
+
+        Args:
+            data_in: Audio input (file path, numpy array, or list).
+            data_lengths: Not used.
+            key (list): Sample identifiers.
+            tokenizer: Not used.
+            frontend: Not used.
+            **kwargs: Must include 'device' (str) and optional 'fs' (int, default 16000).
+
+        Returns:
+            tuple: (results, meta_data) where results is
+                [{"spk_embedding": Tensor of shape (1, 192)}]
+        """
+        # extract fbank feats
+        meta_data = {}
+        time1 = time.perf_counter()
+        audio_sample_list = load_audio_text_image_video(
+            data_in, fs=16000, audio_fs=kwargs.get("fs", 16000), data_type="sound"
+        )
+        time2 = time.perf_counter()
+        meta_data["load_data"] = f"{time2 - time1:0.3f}"
+        speech, speech_lengths, speech_times = extract_feature(audio_sample_list)
+        speech = speech.to(device=kwargs["device"])
+        time3 = time.perf_counter()
+        meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
+        meta_data["batch_data_time"] = np.array(speech_times).sum().item() / 16000.0
+        results = [{"spk_embedding": self.forward(speech.to(torch.float32))}]
+        return results, meta_data

funasr/models/campplus/template.yaml

@@ -0,0 +1,23 @@
+# 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: CAMPPlus
+model_conf:
+    feat_dim: 80
+    embedding_size: 192
+    growth_rate: 32
+    bn_size: 4
+    init_channels: 128
+    config_str: 'batchnorm-relu'
+    memory_efficient: True
+    output_level: 'segment'
+
+# frontend related
+frontend: WavFrontend
+frontend_conf:
+    fs: 16000

funasr/models/campplus/utils.py

@@ -0,0 +1,649 @@
+#!/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)
+# Modified from 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker)
+
+import io
+import os
+import torch
+import requests
+import tempfile
+import contextlib
+import numpy as np
+import librosa as sf
+from typing import Union
+from pathlib import Path
+from typing import Generator, Union
+from abc import ABCMeta, abstractmethod
+import torchaudio.compliance.kaldi as Kaldi
+
+from funasr.models.transformer.utils.nets_utils import pad_list
+
+
+def check_audio_list(audio: list):
+    """Check audio list.
+    
+        Args:
+            audio: TODO.
+        """
+    audio_dur = 0
+    for i in range(len(audio)):
+        seg = audio[i]
+        assert seg[1] >= seg[0], "modelscope error: Wrong time stamps."
+        assert isinstance(seg[2], np.ndarray), "modelscope error: Wrong data type."
+        assert (
+            int(seg[1] * 16000) - int(seg[0] * 16000) == seg[2].shape[0]
+        ), "modelscope error: audio data in list is inconsistent with time length."
+        if i > 0:
+            assert seg[0] >= audio[i - 1][1], "modelscope error: Wrong time stamps."
+        audio_dur += seg[1] - seg[0]
+    return audio_dur
+    # assert audio_dur > 5, 'modelscope error: The effective audio duration is too short.'
+
+
+def sv_preprocess(inputs: Union[np.ndarray, list]):
+    """Sv preprocess.
+    
+        Args:
+            inputs: TODO.
+        """
+    output = []
+    for i in range(len(inputs)):
+        if isinstance(inputs[i], str):
+            file_bytes = File.read(inputs[i])
+            data, fs = sf.load(io.BytesIO(file_bytes), dtype="float32")
+            if len(data.shape) == 2:
+                data = data[:, 0]
+            data = torch.from_numpy(data).unsqueeze(0)
+            data = data.squeeze(0)
+        elif isinstance(inputs[i], np.ndarray):
+            assert len(inputs[i].shape) == 1, "modelscope error: Input array should be [N, T]"
+            data = inputs[i]
+            if data.dtype in ["int16", "int32", "int64"]:
+                data = (data / (1 << 15)).astype("float32")
+            else:
+                data = data.astype("float32")
+            data = torch.from_numpy(data)
+        else:
+            raise ValueError(
+                "modelscope error: The input type is restricted to audio address and nump array."
+            )
+        output.append(data)
+    return output
+
+
+def sv_chunk(vad_segments: list, fs=16000) -> list:
+    """Sv chunk.
+    
+        Args:
+            vad_segments: TODO.
+            fs: TODO.
+        """
+    config = {
+        "seg_dur": 1.5,
+        "seg_shift": 0.75,
+    }
+
+    def seg_chunk(seg_data):
+        """Seg chunk.
+        
+            Args:
+                seg_data: TODO.
+            """
+        seg_st = seg_data[0]
+        data = seg_data[2]
+        chunk_len = int(config["seg_dur"] * fs)
+        chunk_shift = int(config["seg_shift"] * fs)
+        last_chunk_ed = 0
+        seg_res = []
+        for chunk_st in range(0, data.shape[0], chunk_shift):
+            chunk_ed = min(chunk_st + chunk_len, data.shape[0])
+            if chunk_ed <= last_chunk_ed:
+                break
+            last_chunk_ed = chunk_ed
+            chunk_st = max(0, chunk_ed - chunk_len)
+            chunk_data = data[chunk_st:chunk_ed]
+            if chunk_data.shape[0] < chunk_len:
+                chunk_data = np.pad(chunk_data, (0, chunk_len - chunk_data.shape[0]), "constant")
+            seg_res.append([chunk_st / fs + seg_st, chunk_ed / fs + seg_st, chunk_data])
+        return seg_res
+
+    segs = []
+    for i, s in enumerate(vad_segments):
+        segs.extend(seg_chunk(s))
+
+    return segs
+
+
+def extract_feature(audio):
+    """Extract feature.
+    
+        Args:
+            audio: TODO.
+        """
+    features = []
+    feature_times = []
+    feature_lengths = []
+    for au in audio:
+        feature = Kaldi.fbank(au.unsqueeze(0), num_mel_bins=80)
+        feature = feature - feature.mean(dim=0, keepdim=True)
+        features.append(feature)
+        feature_times.append(au.shape[0])
+        feature_lengths.append(feature.shape[0])
+    # padding for batch inference
+    features_padded = pad_list(features, pad_value=0)
+    # features = torch.cat(features)
+    return features_padded, feature_lengths, feature_times
+
+
+def postprocess(
+    segments: list,
+    vad_segments: list,
+    labels: np.ndarray,
+    embeddings: np.ndarray,
+    return_spk_center: bool = False,
+) -> Union[list, tuple]:
+    """Postprocess.
+    
+        Args:
+            segments: TODO.
+            vad_segments: TODO.
+            labels: TODO.
+            embeddings: TODO.
+        """
+    assert len(segments) == len(labels)
+    labels = correct_labels(labels)
+    distribute_res = []
+    for i in range(len(segments)):
+        distribute_res.append([segments[i][0], segments[i][1], labels[i]])
+    # merge the same speakers chronologically
+    distribute_res = merge_seque(distribute_res)
+
+    def is_overlapped(t1, t2):
+        """Is overlapped.
+        
+            Args:
+                t1: TODO.
+                t2: TODO.
+            """
+        if t1 > t2 + 1e-4:
+            return True
+        return False
+
+    # distribute the overlap region
+    for i in range(1, len(distribute_res)):
+        if is_overlapped(distribute_res[i - 1][1], distribute_res[i][0]):
+            p = (distribute_res[i][0] + distribute_res[i - 1][1]) / 2
+            distribute_res[i][0] = p
+            distribute_res[i - 1][1] = p
+
+    # smooth the result
+    distribute_res = smooth(distribute_res)
+
+    if return_spk_center:
+        # spk_embs[i] is the centroid (mean of clustered chunk embeddings) for
+        # corrected speaker label i, aligned with the `spk` ids in sentence_info.
+        # Computed lazily: only when the caller requests speaker centers.
+        spk_embs = np.stack(
+            [embeddings[labels == i].mean(0) for i in range(labels.max() + 1)]
+        )
+        return distribute_res, spk_embs
+    return distribute_res
+
+
+def correct_labels(labels):
+    """Correct labels.
+    
+        Args:
+            labels: TODO.
+        """
+    labels_id = 0
+    id2id = {}
+    new_labels = []
+    for i in labels:
+        if i not in id2id:
+            id2id[i] = labels_id
+            labels_id += 1
+        new_labels.append(id2id[i])
+    return np.array(new_labels)
+
+
+def merge_seque(distribute_res):
+    """Merge seque.
+    
+        Args:
+            distribute_res: TODO.
+        """
+    res = [distribute_res[0]]
+    for i in range(1, len(distribute_res)):
+        if distribute_res[i][2] != res[-1][2] or distribute_res[i][0] > res[-1][1]:
+            res.append(distribute_res[i])
+        else:
+            res[-1][1] = distribute_res[i][1]
+    return res
+
+
+def smooth(res, mindur=0.7):
+    # if only one segment, return directly
+    """Smooth.
+    
+        Args:
+            res: TODO.
+            mindur: TODO.
+        """
+    if len(res) < 2:
+        return res
+    # short segments are assigned to nearest speakers.
+    for i in range(len(res)):
+        res[i][0] = round(res[i][0], 2)
+        res[i][1] = round(res[i][1], 2)
+        if res[i][1] - res[i][0] < mindur:
+            if i == 0:
+                res[i][2] = res[i + 1][2]
+            elif i == len(res) - 1:
+                res[i][2] = res[i - 1][2]
+            elif res[i][0] - res[i - 1][1] <= res[i + 1][0] - res[i][1]:
+                res[i][2] = res[i - 1][2]
+            else:
+                res[i][2] = res[i + 1][2]
+    # merge the speakers
+    res = merge_seque(res)
+
+    return res
+
+
+def distribute_spk(sentence_list, sd_time_list):
+    """Distribute spk.
+    
+        Args:
+            sentence_list: TODO.
+            sd_time_list: TODO.
+        """
+    sd_time_list = [(spk_st * 1000, spk_ed * 1000, spk) for spk_st, spk_ed, spk in sd_time_list]
+    for d in sentence_list:
+        sentence_start = d['start']
+        sentence_end = d['end']
+        sentence_spk = 0
+        max_overlap = 0
+        for spk_st, spk_ed, spk in sd_time_list:
+            overlap = max(min(sentence_end, spk_ed) - max(sentence_start, spk_st), 0)
+            if overlap > max_overlap:
+                max_overlap = overlap
+                sentence_spk = spk
+            if overlap > 0 and sentence_spk == spk:
+                max_overlap += overlap
+        d['spk'] = int(sentence_spk)
+    return sentence_list
+
+
+class Storage(metaclass=ABCMeta):
+    """Abstract class of storage.
+
+    All backends need to implement two apis: ``read()`` and ``read_text()``.
+    ``read()`` reads the file as a byte stream and ``read_text()`` reads
+    the file as texts.
+    """
+
+    @abstractmethod
+    def read(self, filepath: str):
+        """Read.
+        
+            Args:
+                filepath: TODO.
+            """
+        pass
+
+    @abstractmethod
+    def read_text(self, filepath: str):
+        """Read text.
+        
+            Args:
+                filepath: TODO.
+            """
+        pass
+
+    @abstractmethod
+    def write(self, obj: bytes, filepath: Union[str, Path]) -> None:
+        """Write.
+        
+            Args:
+                obj: TODO.
+                filepath: TODO.
+            """
+        pass
+
+    @abstractmethod
+    def write_text(self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8") -> None:
+        """Write text.
+        
+            Args:
+                obj: TODO.
+                filepath: TODO.
+                encoding: TODO.
+            """
+        pass
+
+
+class LocalStorage(Storage):
+    """Local hard disk storage"""
+
+    def read(self, filepath: Union[str, Path]) -> bytes:
+        """Read data from a given ``filepath`` with 'rb' mode.
+
+        Args:
+            filepath (str or Path): Path to read data.
+
+        Returns:
+            bytes: Expected bytes object.
+        """
+        with open(filepath, "rb") as f:
+            content = f.read()
+        return content
+
+    def read_text(self, filepath: Union[str, Path], encoding: str = "utf-8") -> str:
+        """Read data from a given ``filepath`` with 'r' mode.
+
+        Args:
+            filepath (str or Path): Path to read data.
+            encoding (str): The encoding format used to open the ``filepath``.
+                Default: 'utf-8'.
+
+        Returns:
+            str: Expected text reading from ``filepath``.
+        """
+        with open(filepath, "r", encoding=encoding) as f:
+            value_buf = f.read()
+        return value_buf
+
+    def write(self, obj: bytes, filepath: Union[str, Path]) -> None:
+        """Write data to a given ``filepath`` with 'wb' mode.
+
+        Note:
+            ``write`` will create a directory if the directory of ``filepath``
+            does not exist.
+
+        Args:
+            obj (bytes): Data to be written.
+            filepath (str or Path): Path to write data.
+        """
+        dirname = os.path.dirname(filepath)
+        if dirname and not os.path.exists(dirname):
+            os.makedirs(dirname, exist_ok=True)
+
+        with open(filepath, "wb") as f:
+            f.write(obj)
+
+    def write_text(self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8") -> None:
+        """Write data to a given ``filepath`` with 'w' mode.
+
+        Note:
+            ``write_text`` will create a directory if the directory of
+            ``filepath`` does not exist.
+
+        Args:
+            obj (str): Data to be written.
+            filepath (str or Path): Path to write data.
+            encoding (str): The encoding format used to open the ``filepath``.
+                Default: 'utf-8'.
+        """
+        dirname = os.path.dirname(filepath)
+        if dirname and not os.path.exists(dirname):
+            os.makedirs(dirname, exist_ok=True)
+
+        with open(filepath, "w", encoding=encoding) as f:
+            f.write(obj)
+
+    @contextlib.contextmanager
+    def as_local_path(self, filepath: Union[str, Path]) -> Generator[Union[str, Path], None, None]:
+        """Only for unified API and do nothing."""
+        yield filepath
+
+
+class HTTPStorage(Storage):
+    """HTTP and HTTPS storage."""
+
+    def read(self, url):
+        # TODO @wenmeng.zwm add progress bar if file is too large
+        """Read.
+        
+            Args:
+                url: TODO.
+            """
+        r = requests.get(url)
+        r.raise_for_status()
+        return r.content
+
+    def read_text(self, url):
+        """Read text.
+        
+            Args:
+                url: TODO.
+            """
+        r = requests.get(url)
+        r.raise_for_status()
+        return r.text
+
+    @contextlib.contextmanager
+    def as_local_path(self, filepath: str) -> Generator[Union[str, Path], None, None]:
+        """Download a file from ``filepath``.
+
+        ``as_local_path`` is decorated by :meth:`contextlib.contextmanager`. It
+        can be called with ``with`` statement, and when exists from the
+        ``with`` statement, the temporary path will be released.
+
+        Args:
+            filepath (str): Download a file from ``filepath``.
+
+        Examples:
+            >>> storage = HTTPStorage()
+            >>> # After existing from the ``with`` clause,
+            >>> # the path will be removed
+            >>> with storage.get_local_path('http://path/to/file') as path:
+            ...     # do something here
+        """
+        try:
+            f = tempfile.NamedTemporaryFile(delete=False)
+            f.write(self.read(filepath))
+            f.close()
+            yield f.name
+        finally:
+            os.remove(f.name)
+
+    def write(self, obj: bytes, url: Union[str, Path]) -> None:
+        """Write.
+        
+            Args:
+                obj: TODO.
+                url: TODO.
+            """
+        raise NotImplementedError("write is not supported by HTTP Storage")
+
+    def write_text(self, obj: str, url: Union[str, Path], encoding: str = "utf-8") -> None:
+        """Write text.
+        
+            Args:
+                obj: TODO.
+                url: TODO.
+                encoding: TODO.
+            """
+        raise NotImplementedError("write_text is not supported by HTTP Storage")
+
+
+class OSSStorage(Storage):
+    """OSS storage."""
+
+    def __init__(self, oss_config_file=None):
+        # read from config file or env var
+        """Initialize OSSStorage.
+        
+            Args:
+                oss_config_file: TODO.
+            """
+        raise NotImplementedError("OSSStorage.__init__ to be implemented in the future")
+
+    def read(self, filepath):
+        """Read.
+        
+            Args:
+                filepath: TODO.
+            """
+        raise NotImplementedError("OSSStorage.read to be implemented in the future")
+
+    def read_text(self, filepath, encoding="utf-8"):
+        """Read text.
+        
+            Args:
+                filepath: TODO.
+                encoding: TODO.
+            """
+        raise NotImplementedError("OSSStorage.read_text to be implemented in the future")
+
+    @contextlib.contextmanager
+    def as_local_path(self, filepath: str) -> Generator[Union[str, Path], None, None]:
+        """Download a file from ``filepath``.
+
+        ``as_local_path`` is decorated by :meth:`contextlib.contextmanager`. It
+        can be called with ``with`` statement, and when exists from the
+        ``with`` statement, the temporary path will be released.
+
+        Args:
+            filepath (str): Download a file from ``filepath``.
+
+        Examples:
+            >>> storage = OSSStorage()
+            >>> # After existing from the ``with`` clause,
+            >>> # the path will be removed
+            >>> with storage.get_local_path('http://path/to/file') as path:
+            ...     # do something here
+        """
+        try:
+            f = tempfile.NamedTemporaryFile(delete=False)
+            f.write(self.read(filepath))
+            f.close()
+            yield f.name
+        finally:
+            os.remove(f.name)
+
+    def write(self, obj: bytes, filepath: Union[str, Path]) -> None:
+        """Write.
+        
+            Args:
+                obj: TODO.
+                filepath: TODO.
+            """
+        raise NotImplementedError("OSSStorage.write to be implemented in the future")
+
+    def write_text(self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8") -> None:
+        """Write text.
+        
+            Args:
+                obj: TODO.
+                filepath: TODO.
+                encoding: TODO.
+            """
+        raise NotImplementedError("OSSStorage.write_text to be implemented in the future")
+
+
+G_STORAGES = {}
+
+
+class File(object):
+    _prefix_to_storage: dict = {
+        "oss": OSSStorage,
+        "http": HTTPStorage,
+        "https": HTTPStorage,
+        "local": LocalStorage,
+    }
+
+    @staticmethod
+    def _get_storage(uri):
+        """Internal: get storage.
+        
+            Args:
+                uri: TODO.
+            """
+        assert isinstance(uri, str), f"uri should be str type, but got {type(uri)}"
+
+        if "://" not in uri:
+            # local path
+            storage_type = "local"
+        else:
+            prefix, _ = uri.split("://")
+            storage_type = prefix
+
+        assert storage_type in File._prefix_to_storage, (
+            f"Unsupported uri {uri}, valid prefixs: " f"{list(File._prefix_to_storage.keys())}"
+        )
+
+        if storage_type not in G_STORAGES:
+            G_STORAGES[storage_type] = File._prefix_to_storage[storage_type]()
+
+        return G_STORAGES[storage_type]
+
+    @staticmethod
+    def read(uri: str) -> bytes:
+        """Read data from a given ``filepath`` with 'rb' mode.
+
+        Args:
+            filepath (str or Path): Path to read data.
+
+        Returns:
+            bytes: Expected bytes object.
+        """
+        storage = File._get_storage(uri)
+        return storage.read(uri)
+
+    @staticmethod
+    def read_text(uri: Union[str, Path], encoding: str = "utf-8") -> str:
+        """Read data from a given ``filepath`` with 'r' mode.
+
+        Args:
+            filepath (str or Path): Path to read data.
+            encoding (str): The encoding format used to open the ``filepath``.
+                Default: 'utf-8'.
+
+        Returns:
+            str: Expected text reading from ``filepath``.
+        """
+        storage = File._get_storage(uri)
+        return storage.read_text(uri)
+
+    @staticmethod
+    def write(obj: bytes, uri: Union[str, Path]) -> None:
+        """Write data to a given ``filepath`` with 'wb' mode.
+
+        Note:
+            ``write`` will create a directory if the directory of ``filepath``
+            does not exist.
+
+        Args:
+            obj (bytes): Data to be written.
+            filepath (str or Path): Path to write data.
+        """
+        storage = File._get_storage(uri)
+        return storage.write(obj, uri)
+
+    @staticmethod
+    def write_text(obj: str, uri: str, encoding: str = "utf-8") -> None:
+        """Write data to a given ``filepath`` with 'w' mode.
+
+        Note:
+            ``write_text`` will create a directory if the directory of
+            ``filepath`` does not exist.
+
+        Args:
+            obj (str): Data to be written.
+            filepath (str or Path): Path to write data.
+            encoding (str): The encoding format used to open the ``filepath``.
+                Default: 'utf-8'.
+        """
+        storage = File._get_storage(uri)
+        return storage.write_text(obj, uri)
+
+    @contextlib.contextmanager
+    def as_local_path(uri: str) -> Generator[Union[str, Path], None, None]:
+        """Only for unified API and do nothing."""
+        storage = File._get_storage(uri)
+        with storage.as_local_path(uri) as local_path:
+            yield local_path