diff --git a/egs/alimeeting/modular_sa_asr/local/make_textgrid_rttm.py b/egs/alimeeting/modular_sa_asr/local/make_textgrid_rttm.py index f83c572d8..3b6373c78 100755 --- a/egs/alimeeting/modular_sa_asr/local/make_textgrid_rttm.py +++ b/egs/alimeeting/modular_sa_asr/local/make_textgrid_rttm.py @@ -1,7 +1,10 @@ import argparse import tqdm import codecs -import textgrid +try: + import textgrid +except: + raise "Please install textgrid firstly: pip install textgrid" import pdb class Segment(object): diff --git a/egs/alimeeting/modular_sa_asr/local/meeting_speaker_number_process.py b/egs/alimeeting/modular_sa_asr/local/meeting_speaker_number_process.py index 1b09d0af7..8dc98903c 100755 --- a/egs/alimeeting/modular_sa_asr/local/meeting_speaker_number_process.py +++ b/egs/alimeeting/modular_sa_asr/local/meeting_speaker_number_process.py @@ -6,7 +6,10 @@ import argparse import codecs from distutils.util import strtobool from pathlib import Path -import textgrid +try: + import textgrid +except: + raise "Please install textgrid firstly: pip install textgrid" import pdb class Segment(object): diff --git a/egs/alimeeting/sa_asr/local/alimeeting_process_overlap_force.py b/egs/alimeeting/sa_asr/local/alimeeting_process_overlap_force.py index 8ece75706..769003d30 100755 --- a/egs/alimeeting/sa_asr/local/alimeeting_process_overlap_force.py +++ b/egs/alimeeting/sa_asr/local/alimeeting_process_overlap_force.py @@ -6,7 +6,10 @@ import argparse import codecs from distutils.util import strtobool from pathlib import Path -import textgrid +try: + import textgrid +except: + raise "Please install textgrid firstly: pip install textgrid" import pdb class Segment(object): diff --git a/egs/alimeeting/sa_asr/local/alimeeting_process_textgrid.py b/egs/alimeeting/sa_asr/local/alimeeting_process_textgrid.py index 81c19659a..b6d01572e 100755 --- a/egs/alimeeting/sa_asr/local/alimeeting_process_textgrid.py +++ b/egs/alimeeting/sa_asr/local/alimeeting_process_textgrid.py @@ -6,7 +6,10 @@ import argparse import codecs from distutils.util import strtobool from pathlib import Path -import textgrid +try: + import textgrid +except: + raise "Please install textgrid firstly: pip install textgrid" import pdb class Segment(object): diff --git a/egs/alimeeting/sa_asr/local/process_sot_fifo_textchar2spk.py b/egs/alimeeting/sa_asr/local/process_sot_fifo_textchar2spk.py index 488344fb0..c26ba32fe 100755 --- a/egs/alimeeting/sa_asr/local/process_sot_fifo_textchar2spk.py +++ b/egs/alimeeting/sa_asr/local/process_sot_fifo_textchar2spk.py @@ -6,7 +6,10 @@ import argparse import codecs from distutils.util import strtobool from pathlib import Path -import textgrid +try: + import textgrid +except: + raise "Please install textgrid firstly: pip install textgrid" import pdb def get_args(): diff --git a/egs/alimeeting/sa_asr/local/process_textgrid_to_single_speaker_wav.py b/egs/alimeeting/sa_asr/local/process_textgrid_to_single_speaker_wav.py index fdf246090..b72ddc9ef 100755 --- a/egs/alimeeting/sa_asr/local/process_textgrid_to_single_speaker_wav.py +++ b/egs/alimeeting/sa_asr/local/process_textgrid_to_single_speaker_wav.py @@ -6,7 +6,12 @@ import argparse import codecs from distutils.util import strtobool from pathlib import Path -import textgrid + +try: + import textgrid +except: + raise "Please install textgrid firstly: pip install textgrid" + import pdb import numpy as np import sys diff --git a/funasr/bin/asr_infer.py b/funasr/bin/asr_infer.py index 7015eb8e2..c1d08dfc0 100644 --- a/funasr/bin/asr_infer.py +++ b/funasr/bin/asr_infer.py @@ -44,9 +44,9 @@ class Speech2Text: """Speech2Text class Examples: - >>> import soundfile + >>> import librosa >>> speech2text = Speech2Text("asr_config.yml", "asr.pb") - >>> audio, rate = soundfile.read("speech.wav") + >>> audio, rate = librosa.load("speech.wav") >>> speech2text(audio) [(text, token, token_int, hypothesis object), ...] @@ -251,9 +251,9 @@ class Speech2TextParaformer: """Speech2Text class Examples: - >>> import soundfile + >>> import librosa >>> speech2text = Speech2TextParaformer("asr_config.yml", "asr.pb") - >>> audio, rate = soundfile.read("speech.wav") + >>> audio, rate = librosa.load("speech.wav") >>> speech2text(audio) [(text, token, token_int, hypothesis object), ...] @@ -625,9 +625,9 @@ class Speech2TextParaformerOnline: """Speech2Text class Examples: - >>> import soundfile + >>> import librosa >>> speech2text = Speech2TextParaformerOnline("asr_config.yml", "asr.pth") - >>> audio, rate = soundfile.read("speech.wav") + >>> audio, rate = librosa.load("speech.wav") >>> speech2text(audio) [(text, token, token_int, hypothesis object), ...] @@ -876,9 +876,9 @@ class Speech2TextUniASR: """Speech2Text class Examples: - >>> import soundfile + >>> import librosa >>> speech2text = Speech2TextUniASR("asr_config.yml", "asr.pb") - >>> audio, rate = soundfile.read("speech.wav") + >>> audio, rate = librosa.load("speech.wav") >>> speech2text(audio) [(text, token, token_int, hypothesis object), ...] @@ -1106,9 +1106,9 @@ class Speech2TextMFCCA: """Speech2Text class Examples: - >>> import soundfile + >>> import librosa >>> speech2text = Speech2TextMFCCA("asr_config.yml", "asr.pb") - >>> audio, rate = soundfile.read("speech.wav") + >>> audio, rate = librosa.load("speech.wav") >>> speech2text(audio) [(text, token, token_int, hypothesis object), ...] @@ -1637,9 +1637,9 @@ class Speech2TextSAASR: """Speech2Text class Examples: - >>> import soundfile + >>> import librosa >>> speech2text = Speech2TextSAASR("asr_config.yml", "asr.pb") - >>> audio, rate = soundfile.read("speech.wav") + >>> audio, rate = librosa.load("speech.wav") >>> speech2text(audio) [(text, token, token_int, hypothesis object), ...] @@ -1885,9 +1885,9 @@ class Speech2TextWhisper: """Speech2Text class Examples: - >>> import soundfile + >>> import librosa >>> speech2text = Speech2Text("asr_config.yml", "asr.pb") - >>> audio, rate = soundfile.read("speech.wav") + >>> audio, rate = librosa.load("speech.wav") >>> speech2text(audio) [(text, token, token_int, hypothesis object), ...] diff --git a/funasr/bin/asr_inference_launch.py b/funasr/bin/asr_inference_launch.py index e1a32c57c..7dd27fc71 100644 --- a/funasr/bin/asr_inference_launch.py +++ b/funasr/bin/asr_inference_launch.py @@ -20,7 +20,8 @@ from typing import Union import numpy as np import torch import torchaudio -import soundfile +# import librosa +import librosa import yaml from funasr.bin.asr_infer import Speech2Text @@ -1281,7 +1282,8 @@ def inference_paraformer_online( try: raw_inputs = torchaudio.load(data_path_and_name_and_type[0])[0][0] except: - raw_inputs = soundfile.read(data_path_and_name_and_type[0], dtype='float32')[0] + # raw_inputs = librosa.load(data_path_and_name_and_type[0], dtype='float32')[0] + raw_inputs, sr = librosa.load(data_path_and_name_and_type[0], dtype='float32') if raw_inputs.ndim == 2: raw_inputs = raw_inputs[:, 0] raw_inputs = torch.tensor(raw_inputs) diff --git a/funasr/bin/diar_infer.py b/funasr/bin/diar_infer.py index 6fc1da162..bb40f5e41 100755 --- a/funasr/bin/diar_infer.py +++ b/funasr/bin/diar_infer.py @@ -27,11 +27,11 @@ class Speech2DiarizationEEND: """Speech2Diarlization class Examples: - >>> import soundfile + >>> import librosa >>> import numpy as np >>> speech2diar = Speech2DiarizationEEND("diar_sond_config.yml", "diar_sond.pb") >>> profile = np.load("profiles.npy") - >>> audio, rate = soundfile.read("speech.wav") + >>> audio, rate = librosa.load("speech.wav") >>> speech2diar(audio, profile) {"spk1": [(int, int), ...], ...} @@ -109,11 +109,11 @@ class Speech2DiarizationSOND: """Speech2Xvector class Examples: - >>> import soundfile + >>> import librosa >>> import numpy as np >>> speech2diar = Speech2DiarizationSOND("diar_sond_config.yml", "diar_sond.pb") >>> profile = np.load("profiles.npy") - >>> audio, rate = soundfile.read("speech.wav") + >>> audio, rate = librosa.load("speech.wav") >>> speech2diar(audio, profile) {"spk1": [(int, int), ...], ...} diff --git a/funasr/bin/diar_inference_launch.py b/funasr/bin/diar_inference_launch.py index b655df544..f5a11b152 100755 --- a/funasr/bin/diar_inference_launch.py +++ b/funasr/bin/diar_inference_launch.py @@ -15,7 +15,8 @@ from typing import Tuple from typing import Union import numpy as np -import soundfile +# import librosa +import librosa import torch from scipy.signal import medfilt @@ -144,7 +145,9 @@ def inference_sond( # read waveform file example = [load_bytes(x) if isinstance(x, bytes) else x for x in example] - example = [soundfile.read(x)[0] if isinstance(x, str) else x + # example = [librosa.load(x)[0] if isinstance(x, str) else x + # for x in example] + example = [librosa.load(x, dtype='float32')[0] if isinstance(x, str) else x for x in example] # convert torch tensor to numpy array example = [x.numpy() if isinstance(example[0], torch.Tensor) else x diff --git a/funasr/bin/ss_infer.py b/funasr/bin/ss_infer.py index 483967b37..a3eca115d 100644 --- a/funasr/bin/ss_infer.py +++ b/funasr/bin/ss_infer.py @@ -20,9 +20,9 @@ class SpeechSeparator: """SpeechSeparator class Examples: - >>> import soundfile + >>> import librosa >>> speech_separator = MossFormer("ss_config.yml", "ss.pt") - >>> audio, rate = soundfile.read("speech.wav") + >>> audio, rate = librosa.load("speech.wav") >>> separated_wavs = speech_separator(audio) """ diff --git a/funasr/bin/ss_inference_launch.py b/funasr/bin/ss_inference_launch.py index 64503a0a5..0c0241913 100644 --- a/funasr/bin/ss_inference_launch.py +++ b/funasr/bin/ss_inference_launch.py @@ -13,7 +13,7 @@ from typing import Union import numpy as np import torch -import soundfile as sf +import librosa from funasr.build_utils.build_streaming_iterator import build_streaming_iterator from funasr.torch_utils.set_all_random_seed import set_all_random_seed from funasr.utils import config_argparse @@ -104,7 +104,12 @@ def inference_ss( ss_results = speech_separator(**batch) for spk in range(num_spks): - sf.write(os.path.join(output_path, keys[0] + '_s' + str(spk+1)+'.wav'), ss_results[spk], sample_rate) + # sf.write(os.path.join(output_path, keys[0] + '_s' + str(spk+1)+'.wav'), ss_results[spk], sample_rate) + try: + librosa.output.write_wav(os.path.join(output_path, keys[0] + '_s' + str(spk+1)+'.wav'), ss_results[spk], sample_rate) + except: + print("To write wav by librosa, you should install librosa<=0.8.0") + raise torch.cuda.empty_cache() return ss_results diff --git a/funasr/bin/sv_infer.py b/funasr/bin/sv_infer.py index 346440af0..19cfc2e66 100755 --- a/funasr/bin/sv_infer.py +++ b/funasr/bin/sv_infer.py @@ -22,9 +22,9 @@ class Speech2Xvector: """Speech2Xvector class Examples: - >>> import soundfile + >>> import librosa >>> speech2xvector = Speech2Xvector("sv_config.yml", "sv.pb") - >>> audio, rate = soundfile.read("speech.wav") + >>> audio, rate = librosa.load("speech.wav") >>> speech2xvector(audio) [(text, token, token_int, hypothesis object), ...] diff --git a/funasr/bin/vad_infer.py b/funasr/bin/vad_infer.py index 73e1f3f3b..57638739c 100644 --- a/funasr/bin/vad_infer.py +++ b/funasr/bin/vad_infer.py @@ -23,9 +23,9 @@ class Speech2VadSegment: """Speech2VadSegment class Examples: - >>> import soundfile + >>> import librosa >>> speech2segment = Speech2VadSegment("vad_config.yml", "vad.pt") - >>> audio, rate = soundfile.read("speech.wav") + >>> audio, rate = librosa.load("speech.wav") >>> speech2segment(audio) [[10, 230], [245, 450], ...] @@ -118,9 +118,9 @@ class Speech2VadSegmentOnline(Speech2VadSegment): """Speech2VadSegmentOnline class Examples: - >>> import soundfile + >>> import librosa >>> speech2segment = Speech2VadSegmentOnline("vad_config.yml", "vad.pt") - >>> audio, rate = soundfile.read("speech.wav") + >>> audio, rate = librosa.load("speech.wav") >>> speech2segment(audio) [[10, 230], [245, 450], ...] diff --git a/funasr/build_utils/build_trainer.py b/funasr/build_utils/build_trainer.py index 03aa7802d..498d05d4f 100644 --- a/funasr/build_utils/build_trainer.py +++ b/funasr/build_utils/build_trainer.py @@ -246,14 +246,11 @@ class Trainer: for iepoch in range(start_epoch, trainer_options.max_epoch + 1): if iepoch != start_epoch: logging.info( - "{}/{}epoch started. Estimated time to finish: {}".format( + "{}/{}epoch started. Estimated time to finish: {} hours".format( iepoch, trainer_options.max_epoch, - humanfriendly.format_timespan( - (time.perf_counter() - start_time) - / (iepoch - start_epoch) - * (trainer_options.max_epoch - iepoch + 1) - ), + (time.perf_counter() - start_time) / 3600.0 / (iepoch - start_epoch) * ( + trainer_options.max_epoch - iepoch + 1), ) ) else: diff --git a/funasr/datasets/dataset.py b/funasr/datasets/dataset.py index 407f6aa9e..673a9b2c9 100644 --- a/funasr/datasets/dataset.py +++ b/funasr/datasets/dataset.py @@ -16,8 +16,10 @@ from typing import Dict from typing import Mapping from typing import Tuple from typing import Union - -import h5py +try: + import h5py +except: + print("If you want use h5py dataset, please pip install h5py, and try it again") import humanfriendly import kaldiio import numpy as np diff --git a/funasr/datasets/iterable_dataset.py b/funasr/datasets/iterable_dataset.py index 6398e0cf7..b2cc283e1 100644 --- a/funasr/datasets/iterable_dataset.py +++ b/funasr/datasets/iterable_dataset.py @@ -14,7 +14,8 @@ import kaldiio import numpy as np import torch import torchaudio -import soundfile +# import librosa +import librosa from torch.utils.data.dataset import IterableDataset import os.path @@ -70,7 +71,8 @@ def load_wav(input): try: return torchaudio.load(input)[0].numpy() except: - waveform, _ = soundfile.read(input, dtype='float32') + # waveform, _ = librosa.load(input, dtype='float32') + waveform, _ = librosa.load(input, dtype='float32') if waveform.ndim == 2: waveform = waveform[:, 0] return np.expand_dims(waveform, axis=0) diff --git a/funasr/datasets/large_datasets/dataset.py b/funasr/datasets/large_datasets/dataset.py index adfe4f6d9..d3489c199 100644 --- a/funasr/datasets/large_datasets/dataset.py +++ b/funasr/datasets/large_datasets/dataset.py @@ -7,7 +7,8 @@ import torch import torch.distributed as dist import torchaudio import numpy as np -import soundfile +# import librosa +import librosa from kaldiio import ReadHelper from torch.utils.data import IterableDataset @@ -128,7 +129,8 @@ class AudioDataset(IterableDataset): try: waveform, sampling_rate = torchaudio.load(path) except: - waveform, sampling_rate = soundfile.read(path, dtype='float32') + # waveform, sampling_rate = librosa.load(path, dtype='float32') + waveform, sampling_rate = librosa.load(path, dtype='float32') if waveform.ndim == 2: waveform = waveform[:, 0] waveform = np.expand_dims(waveform, axis=0) diff --git a/funasr/datasets/preprocessor.py b/funasr/datasets/preprocessor.py index 9b5c4e77e..26e062c19 100644 --- a/funasr/datasets/preprocessor.py +++ b/funasr/datasets/preprocessor.py @@ -10,7 +10,7 @@ from typing import Union import numpy as np import scipy.signal -import soundfile +import librosa import jieba from funasr.text.build_tokenizer import build_tokenizer @@ -284,7 +284,7 @@ class CommonPreprocessor(AbsPreprocessor): if self.rirs is not None and self.rir_apply_prob >= np.random.random(): rir_path = np.random.choice(self.rirs) if rir_path is not None: - rir, _ = soundfile.read( + rir, _ = librosa.load( rir_path, dtype=np.float64, always_2d=True ) @@ -310,28 +310,31 @@ class CommonPreprocessor(AbsPreprocessor): noise_db = np.random.uniform( self.noise_db_low, self.noise_db_high ) - with soundfile.SoundFile(noise_path) as f: - if f.frames == nsamples: - noise = f.read(dtype=np.float64, always_2d=True) - elif f.frames < nsamples: - offset = np.random.randint(0, nsamples - f.frames) - # noise: (Time, Nmic) - noise = f.read(dtype=np.float64, always_2d=True) - # Repeat noise - noise = np.pad( - noise, - [(offset, nsamples - f.frames - offset), (0, 0)], - mode="wrap", - ) - else: - offset = np.random.randint(0, f.frames - nsamples) - f.seek(offset) - # noise: (Time, Nmic) - noise = f.read( - nsamples, dtype=np.float64, always_2d=True - ) - if len(noise) != nsamples: - raise RuntimeError(f"Something wrong: {noise_path}") + + audio_data = librosa.load(noise_path, dtype='float32')[0][None, :] + frames = len(audio_data[0]) + if frames == nsamples: + noise = audio_data + elif frames < nsamples: + offset = np.random.randint(0, nsamples - frames) + # noise: (Time, Nmic) + noise = audio_data + # Repeat noise + noise = np.pad( + noise, + [(offset, nsamples - frames - offset), (0, 0)], + mode="wrap", + ) + else: + noise = audio_data[:, nsamples] + # offset = np.random.randint(0, frames - nsamples) + # f.seek(offset) + # noise: (Time, Nmic) + # noise = f.read( + # nsamples, dtype=np.float64, always_2d=True + # ) + # if len(noise) != nsamples: + # raise RuntimeError(f"Something wrong: {noise_path}") # noise: (Nmic, Time) noise = noise.T diff --git a/funasr/datasets/small_datasets/preprocessor.py b/funasr/datasets/small_datasets/preprocessor.py index 0ebf325e0..f0d3c9a00 100644 --- a/funasr/datasets/small_datasets/preprocessor.py +++ b/funasr/datasets/small_datasets/preprocessor.py @@ -9,7 +9,7 @@ from typing import Union import numpy as np import scipy.signal -import soundfile +import librosa from funasr.text.build_tokenizer import build_tokenizer from funasr.text.cleaner import TextCleaner @@ -275,7 +275,7 @@ class CommonPreprocessor(AbsPreprocessor): if self.rirs is not None and self.rir_apply_prob >= np.random.random(): rir_path = np.random.choice(self.rirs) if rir_path is not None: - rir, _ = soundfile.read( + rir, _ = librosa.load( rir_path, dtype=np.float64, always_2d=True ) @@ -301,28 +301,30 @@ class CommonPreprocessor(AbsPreprocessor): noise_db = np.random.uniform( self.noise_db_low, self.noise_db_high ) - with soundfile.SoundFile(noise_path) as f: - if f.frames == nsamples: - noise = f.read(dtype=np.float64, always_2d=True) - elif f.frames < nsamples: - offset = np.random.randint(0, nsamples - f.frames) - # noise: (Time, Nmic) - noise = f.read(dtype=np.float64, always_2d=True) - # Repeat noise - noise = np.pad( - noise, - [(offset, nsamples - f.frames - offset), (0, 0)], - mode="wrap", - ) - else: - offset = np.random.randint(0, f.frames - nsamples) - f.seek(offset) - # noise: (Time, Nmic) - noise = f.read( - nsamples, dtype=np.float64, always_2d=True - ) - if len(noise) != nsamples: - raise RuntimeError(f"Something wrong: {noise_path}") + audio_data = librosa.load(noise_path, dtype='float32')[0][None, :] + frames = len(audio_data[0]) + if frames == nsamples: + noise = audio_data + elif frames < nsamples: + offset = np.random.randint(0, nsamples - frames) + # noise: (Time, Nmic) + noise = audio_data + # Repeat noise + noise = np.pad( + noise, + [(offset, nsamples - frames - offset), (0, 0)], + mode="wrap", + ) + else: + noise = audio_data[:, nsamples] + # offset = np.random.randint(0, frames - nsamples) + # f.seek(offset) + # noise: (Time, Nmic) + # noise = f.read( + # nsamples, dtype=np.float64, always_2d=True + # ) + # if len(noise) != nsamples: + # raise RuntimeError(f"Something wrong: {noise_path}") # noise: (Nmic, Time) noise = noise.T diff --git a/funasr/export/export_conformer.py b/funasr/export/export_conformer.py deleted file mode 100644 index 4980775a8..000000000 --- a/funasr/export/export_conformer.py +++ /dev/null @@ -1,151 +0,0 @@ -import json -from typing import Union, Dict -from pathlib import Path - -import os -import logging -import torch - -from funasr.export.models import get_model -import numpy as np -import random -from funasr.utils.types import str2bool, str2triple_str -# torch_version = float(".".join(torch.__version__.split(".")[:2])) -# assert torch_version > 1.9 - -class ModelExport: - def __init__( - self, - cache_dir: Union[Path, str] = None, - onnx: bool = True, - device: str = "cpu", - quant: bool = True, - fallback_num: int = 0, - audio_in: str = None, - calib_num: int = 200, - model_revision: str = None, - ): - self.set_all_random_seed(0) - - self.cache_dir = cache_dir - self.export_config = dict( - feats_dim=560, - onnx=False, - ) - - self.onnx = onnx - self.device = device - self.quant = quant - self.fallback_num = fallback_num - self.frontend = None - self.audio_in = audio_in - self.calib_num = calib_num - self.model_revision = model_revision - - def _export( - self, - model, - model_dir: str = None, - verbose: bool = False, - ): - - export_dir = model_dir - os.makedirs(export_dir, exist_ok=True) - - self.export_config["model_name"] = "model" - model = get_model( - model, - self.export_config, - ) - model.eval() - - if self.onnx: - self._export_onnx(model, verbose, export_dir) - - print("output dir: {}".format(export_dir)) - - def _export_onnx(self, model, verbose, path): - model._export_onnx(verbose, path) - - def set_all_random_seed(self, seed: int): - random.seed(seed) - np.random.seed(seed) - torch.random.manual_seed(seed) - - def parse_audio_in(self, audio_in): - - wav_list, name_list = [], [] - if audio_in.endswith(".scp"): - f = open(audio_in, 'r') - lines = f.readlines()[:self.calib_num] - for line in lines: - name, path = line.strip().split() - name_list.append(name) - wav_list.append(path) - else: - wav_list = [audio_in,] - name_list = ["test",] - return wav_list, name_list - - def load_feats(self, audio_in: str = None): - import torchaudio - - wav_list, name_list = self.parse_audio_in(audio_in) - feats = [] - feats_len = [] - for line in wav_list: - path = line.strip() - waveform, sampling_rate = torchaudio.load(path) - if sampling_rate != self.frontend.fs: - waveform = torchaudio.transforms.Resample(orig_freq=sampling_rate, - new_freq=self.frontend.fs)(waveform) - fbank, fbank_len = self.frontend(waveform, [waveform.size(1)]) - feats.append(fbank) - feats_len.append(fbank_len) - return feats, feats_len - - def export(self, - mode: str = None, - ): - - if mode.startswith('conformer'): - from funasr.tasks.asr import ASRTask - config = os.path.join(model_dir, 'config.yaml') - model_file = os.path.join(model_dir, 'model.pb') - cmvn_file = os.path.join(model_dir, 'am.mvn') - model, asr_train_args = ASRTask.build_model_from_file( - config, model_file, cmvn_file, 'cpu' - ) - self.frontend = model.frontend - self.export_config["feats_dim"] = 560 - - self._export(model, self.cache_dir) - -if __name__ == '__main__': - import argparse - parser = argparse.ArgumentParser() - # parser.add_argument('--model-name', type=str, required=True) - parser.add_argument('--model-name', type=str, action="append", required=True, default=[]) - parser.add_argument('--export-dir', type=str, required=True) - parser.add_argument('--type', type=str, default='onnx', help='["onnx", "torch"]') - 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('--calib_num', type=int, default=200, help='calib max num') - parser.add_argument('--model_revision', type=str, default=None, help='model_revision') - args = parser.parse_args() - - export_model = ModelExport( - cache_dir=args.export_dir, - onnx=args.type == 'onnx', - device=args.device, - quant=args.quantize, - fallback_num=args.fallback_num, - audio_in=args.audio_in, - calib_num=args.calib_num, - model_revision=args.model_revision, - ) - for model_name in args.model_name: - print("export model: {}".format(model_name)) - export_model.export(model_name) diff --git a/funasr/export/models/language_models/__init__.py b/funasr/export/models/language_models/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/funasr/export/models/language_models/embed.py b/funasr/export/models/language_models/embed.py deleted file mode 100644 index 57748f2ea..000000000 --- a/funasr/export/models/language_models/embed.py +++ /dev/null @@ -1,403 +0,0 @@ -"""Positional Encoding Module.""" - -import math - -import torch -import torch.nn as nn -from funasr.modules.embedding import ( - LegacyRelPositionalEncoding, PositionalEncoding, RelPositionalEncoding, - ScaledPositionalEncoding, StreamPositionalEncoding) -from funasr.modules.subsampling import ( - Conv2dSubsampling, Conv2dSubsampling2, Conv2dSubsampling6, - Conv2dSubsampling8) -from funasr.modules.subsampling_without_posenc import \ - Conv2dSubsamplingWOPosEnc - -from funasr.export.models.language_models.subsampling import ( - OnnxConv2dSubsampling, OnnxConv2dSubsampling2, OnnxConv2dSubsampling6, - OnnxConv2dSubsampling8) - - -def get_pos_emb(pos_emb, max_seq_len=512, use_cache=True): - if isinstance(pos_emb, LegacyRelPositionalEncoding): - return OnnxLegacyRelPositionalEncoding(pos_emb, max_seq_len, use_cache) - elif isinstance(pos_emb, ScaledPositionalEncoding): - return OnnxScaledPositionalEncoding(pos_emb, max_seq_len, use_cache) - elif isinstance(pos_emb, RelPositionalEncoding): - return OnnxRelPositionalEncoding(pos_emb, max_seq_len, use_cache) - elif isinstance(pos_emb, PositionalEncoding): - return OnnxPositionalEncoding(pos_emb, max_seq_len, use_cache) - elif isinstance(pos_emb, StreamPositionalEncoding): - return OnnxStreamPositionalEncoding(pos_emb, max_seq_len, use_cache) - elif (isinstance(pos_emb, nn.Sequential) and len(pos_emb) == 0) or ( - isinstance(pos_emb, Conv2dSubsamplingWOPosEnc) - ): - return pos_emb - else: - raise ValueError("Embedding model is not supported.") - - -class Embedding(nn.Module): - def __init__(self, model, max_seq_len=512, use_cache=True): - super().__init__() - self.model = model - if not isinstance(model, nn.Embedding): - if isinstance(model, Conv2dSubsampling): - self.model = OnnxConv2dSubsampling(model) - self.model.out[-1] = get_pos_emb(model.out[-1], max_seq_len) - elif isinstance(model, Conv2dSubsampling2): - self.model = OnnxConv2dSubsampling2(model) - self.model.out[-1] = get_pos_emb(model.out[-1], max_seq_len) - elif isinstance(model, Conv2dSubsampling6): - self.model = OnnxConv2dSubsampling6(model) - self.model.out[-1] = get_pos_emb(model.out[-1], max_seq_len) - elif isinstance(model, Conv2dSubsampling8): - self.model = OnnxConv2dSubsampling8(model) - self.model.out[-1] = get_pos_emb(model.out[-1], max_seq_len) - else: - self.model[-1] = get_pos_emb(model[-1], max_seq_len) - - def forward(self, x, mask=None): - if mask is None: - return self.model(x) - else: - return self.model(x, mask) - - -def _pre_hook( - state_dict, - prefix, - local_metadata, - strict, - missing_keys, - unexpected_keys, - error_msgs, -): - """Perform pre-hook in load_state_dict for backward compatibility. - - Note: - We saved self.pe until v.0.5.2 but we have omitted it later. - Therefore, we remove the item "pe" from `state_dict` for backward compatibility. - - """ - k = prefix + "pe" - if k in state_dict: - state_dict.pop(k) - - -class OnnxPositionalEncoding(torch.nn.Module): - """Positional encoding. - - Args: - d_model (int): Embedding dimension. - dropout_rate (float): Dropout rate. - max_seq_len (int): Maximum input length. - reverse (bool): Whether to reverse the input position. Only for - the class LegacyRelPositionalEncoding. We remove it in the current - class RelPositionalEncoding. - """ - - def __init__(self, model, max_seq_len=512, reverse=False, use_cache=True): - """Construct an PositionalEncoding object.""" - super(OnnxPositionalEncoding, self).__init__() - self.d_model = model.d_model - self.reverse = reverse - self.max_seq_len = max_seq_len - self.xscale = math.sqrt(self.d_model) - self._register_load_state_dict_pre_hook(_pre_hook) - self.pe = model.pe - self.use_cache = use_cache - self.model = model - if self.use_cache: - self.extend_pe() - else: - self.div_term = torch.exp( - torch.arange(0, self.d_model, 2, dtype=torch.float32) - * -(math.log(10000.0) / self.d_model) - ) - - def extend_pe(self): - """Reset the positional encodings.""" - pe_length = len(self.pe[0]) - if self.max_seq_len < pe_length: - self.pe = self.pe[:, : self.max_seq_len] - else: - self.model.extend_pe(torch.tensor(0.0).expand(1, self.max_seq_len)) - self.pe = self.model.pe - - def _add_pe(self, x): - """Computes positional encoding""" - if self.reverse: - position = torch.arange( - x.size(1) - 1, -1, -1.0, dtype=torch.float32 - ).unsqueeze(1) - else: - position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1) - - x = x * self.xscale - x[:, :, 0::2] += torch.sin(position * self.div_term) - x[:, :, 1::2] += torch.cos(position * self.div_term) - return x - - def forward(self, x: torch.Tensor): - """Add positional encoding. - - Args: - x (torch.Tensor): Input tensor (batch, time, `*`). - - Returns: - torch.Tensor: Encoded tensor (batch, time, `*`). - """ - if self.use_cache: - x = x * self.xscale + self.pe[:, : x.size(1)] - else: - x = self._add_pe(x) - return x - - -class OnnxScaledPositionalEncoding(OnnxPositionalEncoding): - """Scaled positional encoding module. - - See Sec. 3.2 https://arxiv.org/abs/1809.08895 - - Args: - d_model (int): Embedding dimension. - dropout_rate (float): Dropout rate. - max_seq_len (int): Maximum input length. - - """ - - def __init__(self, model, max_seq_len=512, use_cache=True): - """Initialize class.""" - super().__init__(model, max_seq_len, use_cache=use_cache) - self.alpha = torch.nn.Parameter(torch.tensor(1.0)) - - def reset_parameters(self): - """Reset parameters.""" - self.alpha.data = torch.tensor(1.0) - - def _add_pe(self, x): - """Computes positional encoding""" - if self.reverse: - position = torch.arange( - x.size(1) - 1, -1, -1.0, dtype=torch.float32 - ).unsqueeze(1) - else: - position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1) - - x = x * self.alpha - x[:, :, 0::2] += torch.sin(position * self.div_term) - x[:, :, 1::2] += torch.cos(position * self.div_term) - return x - - def forward(self, x): - """Add positional encoding. - - Args: - x (torch.Tensor): Input tensor (batch, time, `*`). - - Returns: - torch.Tensor: Encoded tensor (batch, time, `*`). - - """ - if self.use_cache: - x = x + self.alpha * self.pe[:, : x.size(1)] - else: - x = self._add_pe(x) - return x - - -class OnnxLegacyRelPositionalEncoding(OnnxPositionalEncoding): - """Relative positional encoding module (old version). - - Details can be found in https://github.com/espnet/espnet/pull/2816. - - See : Appendix B in https://arxiv.org/abs/1901.02860 - - Args: - d_model (int): Embedding dimension. - dropout_rate (float): Dropout rate. - max_seq_len (int): Maximum input length. - - """ - - def __init__(self, model, max_seq_len=512, use_cache=True): - """Initialize class.""" - super().__init__(model, max_seq_len, reverse=True, use_cache=use_cache) - - def _get_pe(self, x): - """Computes positional encoding""" - if self.reverse: - position = torch.arange( - x.size(1) - 1, -1, -1.0, dtype=torch.float32 - ).unsqueeze(1) - else: - position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1) - - pe = torch.zeros(x.shape) - pe[:, :, 0::2] += torch.sin(position * self.div_term) - pe[:, :, 1::2] += torch.cos(position * self.div_term) - return pe - - def forward(self, x): - """Compute positional encoding. - - Args: - x (torch.Tensor): Input tensor (batch, time, `*`). - - Returns: - torch.Tensor: Encoded tensor (batch, time, `*`). - torch.Tensor: Positional embedding tensor (1, time, `*`). - - """ - x = x * self.xscale - if self.use_cache: - pos_emb = self.pe[:, : x.size(1)] - else: - pos_emb = self._get_pe(x) - return x, pos_emb - - -class OnnxRelPositionalEncoding(torch.nn.Module): - """Relative positional encoding module (new implementation). - Details can be found in https://github.com/espnet/espnet/pull/2816. - See : Appendix B in https://arxiv.org/abs/1901.02860 - Args: - d_model (int): Embedding dimension. - dropout_rate (float): Dropout rate. - max_seq_len (int): Maximum input length. - """ - - def __init__(self, model, max_seq_len=512, use_cache=True): - """Construct an PositionalEncoding object.""" - super(OnnxRelPositionalEncoding, self).__init__() - self.d_model = model.d_model - self.xscale = math.sqrt(self.d_model) - self.pe = None - self.use_cache = use_cache - if self.use_cache: - self.extend_pe(torch.tensor(0.0).expand(1, max_seq_len)) - else: - self.div_term = torch.exp( - torch.arange(0, self.d_model, 2, dtype=torch.float32) - * -(math.log(10000.0) / self.d_model) - ) - - def extend_pe(self, x): - """Reset the positional encodings.""" - if self.pe is not None and self.pe.size(1) >= x.size(1) * 2 - 1: - # self.pe contains both positive and negative parts - # the length of self.pe is 2 * input_len - 1 - if self.pe.dtype != x.dtype or self.pe.device != x.device: - self.pe = self.pe.to(dtype=x.dtype, device=x.device) - return - # Suppose `i` means to the position of query vecotr and `j` means the - # position of key vector. We use position relative positions when keys - # are to the left (i>j) and negative relative positions otherwise (i