diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..600d2d3 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +.vscode \ No newline at end of file diff --git a/Dockerfile.funasr-mr100 b/Dockerfile.funasr-mr100 index 11fd115..34e33f3 100644 --- a/Dockerfile.funasr-mr100 +++ b/Dockerfile.funasr-mr100 @@ -1,20 +1,22 @@ -FROM corex:4.3.0 +FROM git.modelhub.org.cn:9443/enginex-iluvatar/mr-bi150-4.3.0-x86-ubuntu20.04-py3.10-poc-llm-infer:v1.2.3 WORKDIR /root -COPY requirements.txt /root -RUN pip install -r requirements.txt - +RUN sed -i 's|deb.debian.org|archive.debian.org|g' /etc/apt/sources.list \ + && sed -i 's|security.debian.org|archive.debian.org|g' /etc/apt/sources.list \ + && sed -i 's|buster-updates|buster|g' /etc/apt/sources.list \ + && printf 'Acquire::Check-Valid-Until "false";\n' > /etc/apt/apt.conf.d/99no-check-valid \ + && apt-get clean \ + && rm -rf /var/lib/apt/lists/* RUN apt update && apt install -y vim net-tools -RUN pip install funasr==1.2.6 openai-whisper - ADD . /root/ ADD nltk_data.tar.gz /root/ -RUN tar -xvzf nltk_data.tar.gz +ENV NLTK_DATA=/root/nltk_data -RUN cp ./replaced_files/mr_v100/cif_predictor.py /usr/local/lib/python3.10/site-packages/funasr/models/paraformer/ - -EXPOSE 80 -ENTRYPOINT ["bash"] -CMD ["./start_funasr.sh"] \ No newline at end of file +COPY requirements.txt /root +RUN pip install -r /root/requirements.txt -i https://nexus.4pd.io/repository/pypi-all/simple +# Patch files +COPY fastapi_funasr.py /root/fastapi_funasr.py +COPY ./replaced_files/mr_v100/cif_predictor.py /usr/local/lib/python3.10/site-packages/funasr/models/paraformer/ +COPY ./replaced_files/funasr_nano_model.py /usr/local/lib/python3.10/site-packages/funasr/models/fun_asr_nano/model.py diff --git a/fastapi_funasr.py b/fastapi_funasr.py new file mode 100644 index 0000000..a26363b --- /dev/null +++ b/fastapi_funasr.py @@ -0,0 +1,265 @@ +import os +import time +import argparse +import torchaudio +import torch +import traceback + +from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks, Form +import uuid +import uvicorn +from funasr import AutoModel +from funasr.utils.postprocess_utils import rich_transcription_postprocess +from funasr.models.fun_asr_nano.model import FunASRNano + +os.makedirs("./input", exist_ok=True) +status = "Running" +model = None +device = "" +app = FastAPI() + +CUSTOM_DEVICE = os.getenv("CUSTOM_DEVICE", "") +if CUSTOM_DEVICE.startswith("mlu"): + import torch_mlu +elif CUSTOM_DEVICE.startswith("ascend"): + import torch_npu +elif CUSTOM_DEVICE.startswith("pt"): + import torch_dipu + +def make_all_dense(module: torch.nn.Module): + for name, param in list(module.named_parameters(recurse=True)): + if getattr(param, "is_sparse", False) and param.is_sparse: + with torch.no_grad(): + dense = param.to_dense().contiguous() + parent = module + *mods, leaf = name.split(".") + for m in mods: + parent = getattr(parent, m) + setattr(parent, leaf, torch.nn.Parameter(dense, requires_grad=param.requires_grad)) + + # 处理 buffer（如 running_mean 等） + for name, buf in list(module.named_buffers(recurse=True)): + # PyTorch 稀疏张量 layout 不是 strided + if buf.layout != torch.strided: + dense = buf.to_dense().contiguous() + parent = module + *mods, leaf = name.split(".") + for m in mods: + parent = getattr(parent, m) + parent.register_buffer(leaf, dense, persistent=True) + +def split_audio(waveform, sample_rate, segment_seconds=20): + segment_samples = segment_seconds * sample_rate + segments = [] + for i in range(0, waveform.shape[1], segment_samples): + segment = waveform[:, i:i + segment_samples] + if segment.shape[1] > 0: + segments.append(segment) + return segments + +# def determine_model_type(model_name): +# if "sensevoice" in model_name.lower(): +# return "sensevoice" +# elif "whisper" in model_name.lower(): +# return "whisper" +# elif "paraformer" in model_name.lower(): +# return "paraformer" +# elif "conformer" in model_name.lower(): +# return "conformer" +# elif "uniasr" in model_name.lower(): +# return "uni_asr" +# else: +# return "unknown" + +@app.on_event("startup") +def load_model(): + global status, model, device + + config = app.state.config + use_gpu = config.get("use_gpu", True) + model_dir = config.get("model_dir", "/model") + model_type = config.get("model_type", "sensevoice") + warmup = config.get("warmup", False) + print(">> Startup config:") + print(" model_dir =", model_dir, flush=True) + print(" model_type =", model_type, flush=True) + print(" use_gpu =", use_gpu, flush=True) + print(" warmup =", warmup, flush=True) + + device = "cpu" + if use_gpu: + if CUSTOM_DEVICE.startswith("mlu"): + device = "mlu:0" + elif CUSTOM_DEVICE.startswith("ascend"): + device = "npu:1" + else: + device = "cuda:0" + + # 针对加速卡的特殊处理部分 + if device == "cuda:0" and torch.cuda.get_device_name() == "Iluvatar BI-V100" and model_type == "whisper": + # 天垓100情况下的Whisper需要绕过不支持算子 + torch.backends.cuda.enable_flash_sdp(False) + torch.backends.cuda.enable_mem_efficient_sdp(False) + torch.backends.cuda.enable_math_sdp(True) + print(f"device: {device}", flush=True) + + dense_convert = False + if device == "cuda:0" and CUSTOM_DEVICE.startswith("pt") and model_type == "whisper": + dense_convert = True + if device.startswith("npu") and model_type == "whisper": + # Ascend NPU 加载whisper的部分会有Sparse部分device不匹配 + dense_convert = True + + print(f"dense_convert: {dense_convert}", flush=True) + if dense_convert: + model = AutoModel( + model=model_dir, + vad_model=None, + disable_update=True, + device="cpu" + ) + make_all_dense(model.model) + model.model.to(dtype=torch.float32, memory_format=torch.contiguous_format) + model.model.to(device) + model.kwargs["device"] = device + else: + # 不使用VAD, punct，spk模型，就测试原始ASR能力 + model = AutoModel( + model=model_dir, + # vad_model="fsmn-vad", + # vad_kwargs={"max_single_segment_time": 30000}, + vad_model=None, + device=device, + disable_update=True + ) + + if device.startswith("npu") or warmup: + # Ascend NPU由于底层设计的不同，初始化卡的调度比其他卡更复杂，要先进行warmup + print("Start warmup...", flush=True) + res = model.generate(input="warmup.wav") + print("warmup complete.", flush=True) + + status = "Success" + + +def test_funasr(audio_file, lang): + # 推理部分 + waveform, sample_rate = torchaudio.load(audio_file) + # print(waveform.shape) + duration = waveform.shape[1] / sample_rate + segments = split_audio(waveform, sample_rate, segment_seconds=20) + + generated_text = "" + processing_time = 0 + model_type = app.state.config.get("model_type", "sensevoice") + if model_type == "uni_asr": + # uni_asr比较特殊，设计就是处理长音频的（自带VAD），切分的话前20s如果几乎没有人讲话全是音乐直接会报错 + # 因为可能会被切掉所有音频导致实际编解码输入为0 + ts1 = time.time() + res = model.generate( + input=audio_file + ) + generated_text = res[0]["text"] + ts2 = time.time() + processing_time = ts2 - ts1 + else: + # 按照切分的音频依次输入 + for i, segment in enumerate(segments): + segment_path = f"temp_seg_{i}.wav" + torchaudio.save(segment_path, segment, sample_rate) + ts1 = time.time() + if model_type == "sensevoice": + res = model.generate( + input=segment_path, + cache={}, + language="auto", # "zh", "en", "yue", "ja", "ko", "nospeech" + use_itn=True, + batch_size_s=60, + merge_vad=False, + # merge_length_s=15, + ) + text = rich_transcription_postprocess(res[0]["text"]) + elif model_type == "whisper": + DecodingOptions = { + "task": "transcribe", + "language": lang, + "beam_size": None, + "fp16": False, + "without_timestamps": False, + "prompt": None, + } + res = model.generate( + DecodingOptions=DecodingOptions, + input=segment_path, + batch_size_s=0, + ) + text = res[0]["text"] + elif model_type == "paraformer": + res = model.generate( + input=segment_path, + batch_size_s=300 + ) + text = res[0]["text"] + # paraformer模型会一个字一个字输出，中间夹太多空格会影响1-cer的结果 + text = text.replace(" ", "") + elif model_type == "conformer": + res = model.generate( + input=segment_path, + batch_size_s=300 + ) + text = res[0]["text"] + # elif model_type == "uni_asr": + # if i == 0: + # os.remove(segment_path) + # continue + # res = model.generate( + # input=segment_path + # ) + # text = res[0]["text"] + else: + raise RuntimeError("unknown model type") + ts2 = time.time() + generated_text += text + processing_time += (ts2 - ts1) + os.remove(segment_path) + + rtf = processing_time / duration + print("Text:", generated_text, flush=True) + print(f"Audio duration:\t{duration:.3f} s", flush=True) + print(f"Elapsed:\t{processing_time:.3f} s", flush=True) + print(f"RTF = {processing_time:.3f}/{duration:.3f} = {rtf:.3f}", flush=True) + + return generated_text + +@app.get("/health") +def health(): + + if status=="Running": + return { + "status":"loading model" + } + ret = { + "status": "ok" if status == "Success" else "failed", + } + return ret + +@app.post("/transduce") +def transduce( + audio: UploadFile = File(...), + lang: str = Form("zh"), + background_tasks: BackgroundTasks = None +): + try: + file_path = f"./input/{uuid.uuid4()}.wav" + with open(file_path, "wb") as f: + f.write(audio.file.read()) + background_tasks.add_task(os.remove, file_path) + generated_text = test_funasr(file_path, lang) + + return {"generated_text": generated_text} + except Exception: + raise HTTPException(status_code=500, detail=f"Processing failed: \n{traceback.format_exc()}") + +# if __name__ == "__main__": + +# uvicorn.run("fastapi_funasr:app", host="0.0.0.0", port=1111, workers=1) diff --git a/lei-jun-test.wav b/lei-jun-test.wav deleted file mode 100644 index 84f4482..0000000 Binary files a/lei-jun-test.wav and /dev/null differ diff --git a/lei-jun.txt b/lei-jun.txt deleted file mode 100644 index fc8a560..0000000 --- a/lei-jun.txt +++ /dev/null @@ -1 +0,0 @@ -朋友们晚上好，欢迎大家来参加今天晚上的活动，谢谢大家。这是我第四次办年度演讲，前三次呢因为疫情的原因，都在小米科技园内举办。现场呢人很少。这是第四次，我们仔细想了想，我们还是想办一个比较大的聚会。然后呢让我们的新朋友老朋友一起聚一聚。今天的话呢我们就在北京的国家会议中心呢举办了这么一个活动。现场呢来了很多人大概有3500人。还有很多很多的朋友呢，通过观看直播的方式来参与。再一次呢对大家的参加，表示感谢，谢谢大家。两个月前我参加了今年武汉大学的毕业典礼。今年呢是武汉大学建校130周年，作为校友被母校邀请在毕业典礼上致辞，这对我来说，是至高无上的荣誉。站在讲台的那一刻，面对全校师生，关于武大的所有的记忆一下子涌现在脑海里。今天呢我就先和大家聊聊武大往事。那还是36年前，1987年我呢考上了武汉大学的计算机系。在武汉大学的图书馆里看了一本书《硅谷之火》，建立了我一生的梦想。看完书以后，热血沸腾，激动得睡不着觉。我还记得那天晚上星光很亮。我就在武大的操场上，就是屏幕上这个操场，走了一圈又一圈，走了整整一个晚上。我心里有团火，我也想办一个伟大的公司，就是这样，梦想之火，在我心里彻底点燃了。但是一个大一的新生，但是一个大一的新生，一个从县城里出来的年轻人，什么也不会，什么也没有，就想创办一家伟大的公司，这不就是天方夜谭吗？这么离谱的一个梦想，该如何实现呢？那天晚上我想了一整晚上，说实话，越想越糊涂，完全理不清头绪，后来我在想，哎干脆别想了，把书念好是正事，所以呢我就下定决心认认真真读书，那么我怎么能够把书读得不同凡响呢？ \ No newline at end of file diff --git a/main.py b/main.py new file mode 100644 index 0000000..ae902f6 --- /dev/null +++ b/main.py @@ -0,0 +1,27 @@ +import argparse +import uvicorn +from fastapi_funasr import app + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--model_dir", type=str, default="/model", help="model directory") + parser.add_argument("--model_type", type=str, default="sensevoice", help="model type, e.g. sensevoice") + parser.add_argument("--use_gpu", action="store_true", default=True) + parser.add_argument("--warmup", action="store_true", help="whether do warmup when first initializing model") + parser.add_argument("--port", type=int, default=8000, help="service port") + + args = parser.parse_args() + + # 将参数加到 app.state 中 + app.state.config = { + "model_dir": args.model_dir, + "model_type": args.model_type, + "use_gpu": args.use_gpu, # True + "warmup": args.warmup + } + + uvicorn.run("fastapi_funasr:app", + host="0.0.0.0", + port=args.port, + workers=1 + ) \ No newline at end of file diff --git a/nltk_data.tar.gz b/nltk_data.tar.gz index 68a9b4c..0ccae52 100644 Binary files a/nltk_data.tar.gz and b/nltk_data.tar.gz differ diff --git a/replaced_files/bi_v100/cif_predictor.py b/replaced_files/bi_v100/cif_predictor.py new file mode 100644 index 0000000..db9d1a9 --- /dev/null +++ b/replaced_files/bi_v100/cif_predictor.py @@ -0,0 +1,762 @@ +#!/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 logging +import numpy as np + +from funasr.register import tables +from funasr.train_utils.device_funcs import to_device +from funasr.models.transformer.utils.nets_utils import make_pad_mask +from torch.cuda.amp import autocast + + +@tables.register("predictor_classes", "CifPredictor") +class CifPredictor(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.45, + ): + super().__init__() + + self.pad = torch.nn.ConstantPad1d((l_order, r_order), 0) + self.cif_conv1d = torch.nn.Conv1d(idim, idim, l_order + r_order + 1, groups=idim) + 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 + + def forward( + self, + hidden, + target_label=None, + mask=None, + ignore_id=-1, + mask_chunk_predictor=None, + target_label_length=None, + ): + + with autocast(False): + h = hidden + context = h.transpose(1, 2) + queries = self.pad(context) + memory = self.cif_conv1d(queries) + output = memory + context + output = self.dropout(output) + output = output.transpose(1, 2) + output = torch.relu(output) + 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 + + def tail_process_fn(self, hidden, alphas, token_num=None, mask=None): + 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 + ): + 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", "CifPredictorV2") +class CifPredictorV2(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", + tail_mask=True, + ): + super().__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.tail_mask = tail_mask + + def forward( + self, + hidden, + target_label=None, + mask=None, + ignore_id=-1, + mask_chunk_predictor=None, + target_label_length=None, + ): + + with autocast(False): + 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) + 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.squeeze(-1) + 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: + if self.tail_mask: + hidden, alphas, token_num = self.tail_process_fn( + hidden, alphas, token_num, mask=mask + ) + else: + hidden, alphas, token_num = self.tail_process_fn( + hidden, alphas, token_num, mask=None + ) + acoustic_embeds, cif_peak = cif_v1(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 + + def forward_chunk(self, hidden, cache=None, **kwargs): + is_final = kwargs.get("is_final", False) + batch_size, len_time, hidden_size = hidden.shape + 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) + + alphas = alphas.squeeze(-1) + + token_length = [] + list_fires = [] + list_frames = [] + cache_alphas = [] + cache_hiddens = [] + + if cache is not None and "chunk_size" in cache: + alphas[:, : cache["chunk_size"][0]] = 0.0 + if not is_final: + alphas[:, sum(cache["chunk_size"][:2]) :] = 0.0 + if cache is not None and "cif_alphas" in cache and "cif_hidden" in cache: + cache["cif_hidden"] = to_device(cache["cif_hidden"], device=hidden.device) + cache["cif_alphas"] = to_device(cache["cif_alphas"], device=alphas.device) + hidden = torch.cat((cache["cif_hidden"], hidden), dim=1) + alphas = torch.cat((cache["cif_alphas"], alphas), dim=1) + if cache is not None and is_final: + tail_hidden = torch.zeros((batch_size, 1, hidden_size), device=hidden.device) + tail_alphas = torch.tensor([[self.tail_threshold]], device=alphas.device) + tail_alphas = torch.tile(tail_alphas, (batch_size, 1)) + hidden = torch.cat((hidden, tail_hidden), dim=1) + alphas = torch.cat((alphas, tail_alphas), dim=1) + + len_time = alphas.shape[1] + for b in range(batch_size): + integrate = 0.0 + frames = torch.zeros((hidden_size), device=hidden.device) + list_frame = [] + list_fire = [] + for t in range(len_time): + alpha = alphas[b][t] + if alpha + integrate < self.threshold: + integrate += alpha + list_fire.append(integrate) + frames += alpha * hidden[b][t] + else: + frames += (self.threshold - integrate) * hidden[b][t] + list_frame.append(frames) + integrate += alpha + list_fire.append(integrate) + integrate -= self.threshold + frames = integrate * hidden[b][t] + + cache_alphas.append(integrate) + if integrate > 0.0: + cache_hiddens.append(frames / integrate) + else: + cache_hiddens.append(frames) + + token_length.append(torch.tensor(len(list_frame), device=alphas.device)) + list_fires.append(list_fire) + list_frames.append(list_frame) + + cache["cif_alphas"] = torch.stack(cache_alphas, axis=0) + cache["cif_alphas"] = torch.unsqueeze(cache["cif_alphas"], axis=0) + cache["cif_hidden"] = torch.stack(cache_hiddens, axis=0) + cache["cif_hidden"] = torch.unsqueeze(cache["cif_hidden"], axis=0) + + max_token_len = max(token_length) + if max_token_len == 0: + return hidden, torch.stack(token_length, 0), None, None + list_ls = [] + for b in range(batch_size): + pad_frames = torch.zeros( + (max_token_len - token_length[b], hidden_size), device=alphas.device + ) + if token_length[b] == 0: + list_ls.append(pad_frames) + else: + list_frames[b] = torch.stack(list_frames[b]) + list_ls.append(torch.cat((list_frames[b], pad_frames), dim=0)) + + cache["cif_alphas"] = torch.stack(cache_alphas, axis=0) + cache["cif_alphas"] = torch.unsqueeze(cache["cif_alphas"], axis=0) + cache["cif_hidden"] = torch.stack(cache_hiddens, axis=0) + cache["cif_hidden"] = torch.unsqueeze(cache["cif_hidden"], axis=0) + return torch.stack(list_ls, 0), torch.stack(token_length, 0), None, None + + def tail_process_fn(self, hidden, alphas, token_num=None, mask=None): + 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)) + if b > 1: + alphas = torch.cat([alphas, tail_threshold.repeat(b, 1)], dim=1) + else: + 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 + ): + 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", "CifPredictorV2Export") +class CifPredictorV2Export(torch.nn.Module): + def __init__(self, model, **kwargs): + 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 + + def forward( + self, + hidden: torch.Tensor, + mask: torch.Tensor, + ): + alphas, token_num = self.forward_cnn(hidden, mask) + mask = mask.transpose(-1, -2).float() + mask = mask.squeeze(-1) + hidden, alphas, token_num = self.tail_process_fn(hidden, alphas, mask=mask) + acoustic_embeds, cif_peak = cif_v1_export(hidden, alphas, self.threshold) + + return acoustic_embeds, token_num, alphas, cif_peak + + def forward_cnn( + self, + hidden: torch.Tensor, + mask: torch.Tensor, + ): + 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) + + return alphas, token_num + + def tail_process_fn(self, hidden, alphas, token_num=None, mask=None): + 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_v1_export(hidden, alphas, threshold: float): + device = hidden.device + dtype = hidden.dtype + batch_size, len_time, hidden_size = hidden.size() + threshold = torch.tensor([threshold], dtype=alphas.dtype).to(alphas.device) + + frames = torch.zeros(batch_size, len_time, hidden_size, dtype=dtype, device=device) + fires = torch.zeros(batch_size, len_time, dtype=dtype, device=device) + + # prefix_sum = torch.cumsum(alphas, dim=1) + prefix_sum = torch.cumsum(alphas, dim=1, dtype=torch.float64).to( + torch.float32 + ) # cumsum precision degradation cause wrong result in extreme + prefix_sum_floor = torch.floor(prefix_sum) + dislocation_prefix_sum = torch.roll(prefix_sum, 1, dims=1) + dislocation_prefix_sum_floor = torch.floor(dislocation_prefix_sum) + + dislocation_prefix_sum_floor[:, 0] = 0 + dislocation_diff = prefix_sum_floor - dislocation_prefix_sum_floor + + fire_idxs = dislocation_diff > 0 + fires[fire_idxs] = 1 + fires = fires + prefix_sum - prefix_sum_floor + + # prefix_sum_hidden = torch.cumsum(alphas.unsqueeze(-1).tile((1, 1, hidden_size)) * hidden, dim=1) + prefix_sum_hidden = torch.cumsum(alphas.unsqueeze(-1).repeat((1, 1, hidden_size)) * hidden, dim=1) + frames = prefix_sum_hidden[fire_idxs] + shift_frames = torch.roll(frames, 1, dims=0) + + batch_len = fire_idxs.sum(1) + batch_idxs = torch.cumsum(batch_len, dim=0) + shift_batch_idxs = torch.roll(batch_idxs, 1, dims=0) + shift_batch_idxs[0] = 0 + shift_frames[shift_batch_idxs] = 0 + + remains = fires - torch.floor(fires) + # remain_frames = remains[fire_idxs].unsqueeze(-1).tile((1, hidden_size)) * hidden[fire_idxs] + remain_frames = remains[fire_idxs].unsqueeze(-1).repeat((1, hidden_size)) * hidden[fire_idxs] + + shift_remain_frames = torch.roll(remain_frames, 1, dims=0) + shift_remain_frames[shift_batch_idxs] = 0 + + frames = frames - shift_frames + shift_remain_frames - remain_frames + + # max_label_len = batch_len.max() + max_label_len = alphas.sum(dim=-1) + max_label_len = torch.floor(max_label_len).max().to(dtype=torch.int64) + + # frame_fires = torch.zeros(batch_size, max_label_len, hidden_size, dtype=dtype, device=device) + frame_fires = torch.zeros(batch_size, max_label_len, hidden_size, dtype=dtype, device=device) + indices = torch.arange(max_label_len, device=device).expand(batch_size, -1) + frame_fires_idxs = indices < batch_len.unsqueeze(1) + frame_fires[frame_fires_idxs] = frames + return frame_fires, fires + + +@torch.jit.script +def cif_export(hidden, alphas, threshold: float): + 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 + + +class mae_loss(torch.nn.Module): + + def __init__(self, normalize_length=False): + 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): + 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): + 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()) + 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_v1(alphas, threshold, return_fire_idxs=False): + batch_size, len_time = alphas.size() + device = alphas.device + dtype = alphas.dtype + + threshold = torch.tensor([threshold], dtype=alphas.dtype).to(alphas.device) + + fires = torch.zeros(batch_size, len_time, dtype=dtype, device=device) + + if torch.cuda.get_device_name() == "Iluvatar BI-V100": + # the normal branch causes wrong result in bi-100, and leads to exception in later stages + prefix_sum = torch.cumsum(alphas, dim=1) + else: + prefix_sum = torch.cumsum(alphas, dim=1, dtype=torch.float64).to( + torch.float32 + ) # cumsum precision degradation cause wrong result in extreme + prefix_sum_floor = torch.floor(prefix_sum) + dislocation_prefix_sum = torch.roll(prefix_sum, 1, dims=1) + dislocation_prefix_sum_floor = torch.floor(dislocation_prefix_sum) + + dislocation_prefix_sum_floor[:, 0] = 0 + dislocation_diff = prefix_sum_floor - dislocation_prefix_sum_floor + + fire_idxs = dislocation_diff > 0 + fires[fire_idxs] = 1 + fires = fires + prefix_sum - prefix_sum_floor + if return_fire_idxs: + return fires, fire_idxs + return fires + + +def cif_v1(hidden, alphas, threshold): + fires, fire_idxs = cif_wo_hidden_v1(alphas, threshold, return_fire_idxs=True) + + device = hidden.device + dtype = hidden.dtype + batch_size, len_time, hidden_size = hidden.size() + # frames = torch.zeros(batch_size, len_time, hidden_size, dtype=dtype, device=device) + # prefix_sum_hidden = torch.cumsum(alphas.unsqueeze(-1).tile((1, 1, hidden_size)) * hidden, dim=1) + frames = torch.zeros(batch_size, len_time, hidden_size, dtype=dtype, device=device) + prefix_sum_hidden = torch.cumsum(alphas.unsqueeze(-1).repeat((1, 1, hidden_size)) * hidden, dim=1) + + frames = prefix_sum_hidden[fire_idxs] + shift_frames = torch.roll(frames, 1, dims=0) + + batch_len = fire_idxs.sum(1) + batch_idxs = torch.cumsum(batch_len, dim=0) + shift_batch_idxs = torch.roll(batch_idxs, 1, dims=0) + shift_batch_idxs[0] = 0 + shift_frames[shift_batch_idxs] = 0 + + remains = fires - torch.floor(fires) + # remain_frames = remains[fire_idxs].unsqueeze(-1).tile((1, hidden_size)) * hidden[fire_idxs] + remain_frames = remains[fire_idxs].unsqueeze(-1).repeat((1, hidden_size)) * hidden[fire_idxs] + + shift_remain_frames = torch.roll(remain_frames, 1, dims=0) + shift_remain_frames[shift_batch_idxs] = 0 + + frames = frames - shift_frames + shift_remain_frames - remain_frames + + # max_label_len = batch_len.max() + max_label_len = ( + torch.round(alphas.sum(-1)).int().max() + ) # torch.round to calculate the max length + + # frame_fires = torch.zeros(batch_size, max_label_len, hidden_size, dtype=dtype, device=device) + frame_fires = torch.zeros(batch_size, max_label_len, hidden_size, dtype=dtype, device=device) + indices = torch.arange(max_label_len, device=device).expand(batch_size, -1) + frame_fires_idxs = indices < batch_len.unsqueeze(1) + frame_fires[frame_fires_idxs] = frames + return frame_fires, fires + + +def cif_wo_hidden(alphas, threshold): + 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 diff --git a/replaced_files/bi_v100/model.py b/replaced_files/bi_v100/model.py new file mode 100644 index 0000000..c95e881 --- /dev/null +++ b/replaced_files/bi_v100/model.py @@ -0,0 +1,1032 @@ +#!/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 torch.cuda.amp import autocast +from typing import Union, Dict, List, Tuple, Optional + +from funasr.register import tables +from funasr.models.ctc.ctc import CTC +from funasr.utils import postprocess_utils +from funasr.metrics.compute_acc import th_accuracy +from funasr.utils.datadir_writer import DatadirWriter +from funasr.models.paraformer.cif_predictor import mae_loss +from funasr.train_utils.device_funcs import force_gatherable +from funasr.losses.label_smoothing_loss import LabelSmoothingLoss +from funasr.models.transformer.utils.add_sos_eos import add_sos_eos +from funasr.models.transformer.utils.nets_utils import make_pad_mask, pad_list +from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank +from funasr.models.scama.utils import sequence_mask + + +@tables.register("model_classes", "UniASR") +class UniASR(torch.nn.Module): + """ + Author: Speech Lab of DAMO Academy, Alibaba Group + """ + + def __init__( + self, + specaug: str = None, + specaug_conf: dict = None, + normalize: str = None, + normalize_conf: dict = None, + encoder: str = None, + encoder_conf: dict = None, + encoder2: str = None, + encoder2_conf: dict = None, + decoder: str = None, + decoder_conf: dict = None, + decoder2: str = None, + decoder2_conf: dict = None, + predictor: str = None, + predictor_conf: dict = None, + predictor_bias: int = 0, + predictor_weight: float = 0.0, + predictor2: str = None, + predictor2_conf: dict = None, + predictor2_bias: int = 0, + predictor2_weight: float = 0.0, + ctc: str = None, + ctc_conf: dict = None, + ctc_weight: float = 0.5, + ctc2: str = None, + ctc2_conf: dict = None, + ctc2_weight: float = 0.5, + decoder_attention_chunk_type: str = "chunk", + decoder_attention_chunk_type2: str = "chunk", + stride_conv=None, + stride_conv_conf: dict = None, + loss_weight_model1: float = 0.5, + input_size: int = 80, + vocab_size: int = -1, + ignore_id: int = -1, + blank_id: int = 0, + sos: int = 1, + eos: int = 2, + lsm_weight: float = 0.0, + length_normalized_loss: bool = False, + share_embedding: bool = False, + **kwargs, + ): + super().__init__() + + if specaug is not None: + specaug_class = tables.specaug_classes.get(specaug) + specaug = specaug_class(**specaug_conf) + if normalize is not None: + normalize_class = tables.normalize_classes.get(normalize) + normalize = normalize_class(**normalize_conf) + + encoder_class = tables.encoder_classes.get(encoder) + encoder = encoder_class(input_size=input_size, **encoder_conf) + encoder_output_size = encoder.output_size() + + decoder_class = tables.decoder_classes.get(decoder) + decoder = decoder_class( + vocab_size=vocab_size, + encoder_output_size=encoder_output_size, + **decoder_conf, + ) + predictor_class = tables.predictor_classes.get(predictor) + predictor = predictor_class(**predictor_conf) + + from funasr.models.transformer.utils.subsampling import Conv1dSubsampling + + stride_conv = Conv1dSubsampling( + **stride_conv_conf, + idim=input_size + encoder_output_size, + odim=input_size + encoder_output_size, + ) + stride_conv_output_size = stride_conv.output_size() + + encoder_class = tables.encoder_classes.get(encoder2) + encoder2 = encoder_class(input_size=stride_conv_output_size, **encoder2_conf) + encoder2_output_size = encoder2.output_size() + + decoder_class = tables.decoder_classes.get(decoder2) + decoder2 = decoder_class( + vocab_size=vocab_size, + encoder_output_size=encoder2_output_size, + **decoder2_conf, + ) + predictor_class = tables.predictor_classes.get(predictor2) + predictor2 = predictor_class(**predictor2_conf) + + self.blank_id = blank_id + self.sos = sos + self.eos = eos + self.vocab_size = vocab_size + self.ignore_id = ignore_id + self.ctc_weight = ctc_weight + self.ctc2_weight = ctc2_weight + + self.specaug = specaug + self.normalize = normalize + + self.encoder = encoder + + self.error_calculator = None + + self.decoder = decoder + self.ctc = None + self.ctc2 = None + + self.criterion_att = LabelSmoothingLoss( + size=vocab_size, + padding_idx=ignore_id, + smoothing=lsm_weight, + normalize_length=length_normalized_loss, + ) + + self.predictor = predictor + self.predictor_weight = predictor_weight + self.criterion_pre = mae_loss(normalize_length=length_normalized_loss) + self.encoder1_encoder2_joint_training = kwargs.get("encoder1_encoder2_joint_training", True) + + if self.encoder.overlap_chunk_cls is not None: + from funasr.models.scama.chunk_utilis import ( + build_scama_mask_for_cross_attention_decoder, + ) + + self.build_scama_mask_for_cross_attention_decoder_fn = ( + build_scama_mask_for_cross_attention_decoder + ) + self.decoder_attention_chunk_type = decoder_attention_chunk_type + + self.encoder2 = encoder2 + self.decoder2 = decoder2 + self.ctc2_weight = ctc2_weight + + self.predictor2 = predictor2 + self.predictor2_weight = predictor2_weight + self.decoder_attention_chunk_type2 = decoder_attention_chunk_type2 + self.stride_conv = stride_conv + self.loss_weight_model1 = loss_weight_model1 + if self.encoder2.overlap_chunk_cls is not None: + from funasr.models.scama.chunk_utilis import ( + build_scama_mask_for_cross_attention_decoder, + ) + + self.build_scama_mask_for_cross_attention_decoder_fn2 = ( + build_scama_mask_for_cross_attention_decoder + ) + self.decoder_attention_chunk_type2 = decoder_attention_chunk_type2 + + self.length_normalized_loss = length_normalized_loss + self.enable_maas_finetune = kwargs.get("enable_maas_finetune", False) + self.freeze_encoder2 = kwargs.get("freeze_encoder2", False) + self.beam_search = None + + def forward( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + text: torch.Tensor, + text_lengths: torch.Tensor, + **kwargs, + ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]: + """Frontend + Encoder + Decoder + Calc loss + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + text: (Batch, Length) + text_lengths: (Batch,) + """ + decoding_ind = kwargs.get("decoding_ind", None) + 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] + + ind = self.encoder.overlap_chunk_cls.random_choice(self.training, decoding_ind) + # 1. Encoder + if self.enable_maas_finetune: + with torch.no_grad(): + speech_raw, encoder_out, encoder_out_lens = self.encode( + speech, speech_lengths, ind=ind + ) + else: + speech_raw, encoder_out, encoder_out_lens = self.encode(speech, speech_lengths, ind=ind) + + loss_att, acc_att, cer_att, wer_att = None, None, None, None + loss_ctc, cer_ctc = None, None + stats = dict() + loss_pre = None + loss, loss1, loss2 = 0.0, 0.0, 0.0 + + if self.loss_weight_model1 > 0.0: + ## model1 + # 1. CTC branch + if self.enable_maas_finetune: + with torch.no_grad(): + + loss_att, acc_att, cer_att, wer_att, loss_pre = self._calc_att_predictor_loss( + encoder_out, encoder_out_lens, text, text_lengths + ) + + loss = loss_att + loss_pre * self.predictor_weight + + # Collect Attn branch stats + stats["loss_att"] = loss_att.detach() if loss_att is not None else None + stats["acc"] = acc_att + stats["cer"] = cer_att + stats["wer"] = wer_att + stats["loss_pre"] = loss_pre.detach().cpu() if loss_pre is not None else None + else: + + loss_att, acc_att, cer_att, wer_att, loss_pre = self._calc_att_predictor_loss( + encoder_out, encoder_out_lens, text, text_lengths + ) + + loss = loss_att + loss_pre * self.predictor_weight + + # Collect Attn branch stats + stats["loss_att"] = loss_att.detach() if loss_att is not None else None + stats["acc"] = acc_att + stats["cer"] = cer_att + stats["wer"] = wer_att + stats["loss_pre"] = loss_pre.detach().cpu() if loss_pre is not None else None + + loss1 = loss + + if self.loss_weight_model1 < 1.0: + ## model2 + + # encoder2 + if self.freeze_encoder2: + with torch.no_grad(): + encoder_out, encoder_out_lens = self.encode2( + encoder_out, encoder_out_lens, speech_raw, speech_lengths, ind=ind + ) + else: + encoder_out, encoder_out_lens = self.encode2( + encoder_out, encoder_out_lens, speech_raw, speech_lengths, ind=ind + ) + + intermediate_outs = None + if isinstance(encoder_out, tuple): + intermediate_outs = encoder_out[1] + encoder_out = encoder_out[0] + + loss_att, acc_att, cer_att, wer_att, loss_pre = self._calc_att_predictor_loss2( + encoder_out, encoder_out_lens, text, text_lengths + ) + + loss = loss_att + loss_pre * self.predictor2_weight + + # Collect Attn branch stats + stats["loss_att2"] = loss_att.detach() if loss_att is not None else None + stats["acc2"] = acc_att + stats["cer2"] = cer_att + stats["wer2"] = wer_att + stats["loss_pre2"] = loss_pre.detach().cpu() if loss_pre is not None else None + + loss2 = loss + + loss = loss1 * self.loss_weight_model1 + loss2 * (1 - self.loss_weight_model1) + stats["loss1"] = torch.clone(loss1.detach()) + stats["loss2"] = torch.clone(loss2.detach()) + 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 + 1).sum()) + + loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) + return loss, stats, weight + + def collect_feats( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + text: torch.Tensor, + text_lengths: torch.Tensor, + ) -> Dict[str, torch.Tensor]: + if self.extract_feats_in_collect_stats: + feats, feats_lengths = self._extract_feats(speech, speech_lengths) + else: + # Generate dummy stats if extract_feats_in_collect_stats is False + logging.warning( + "Generating dummy stats for feats and feats_lengths, " + "because encoder_conf.extract_feats_in_collect_stats is " + f"{self.extract_feats_in_collect_stats}" + ) + feats, feats_lengths = speech, speech_lengths + return {"feats": feats, "feats_lengths": feats_lengths} + + def encode( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + **kwargs, + ): + """Frontend + Encoder. Note that this method is used by asr_inference.py + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + """ + ind = kwargs.get("ind", 0) + with autocast(False): + # Data augmentation + if self.specaug is not None and self.training: + speech, speech_lengths = self.specaug(speech, speech_lengths) + + # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN + if self.normalize is not None: + speech, speech_lengths = self.normalize(speech, speech_lengths) + + speech_raw = speech.clone().to(speech.device) + + # 4. Forward encoder + encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths, ind=ind) + if isinstance(encoder_out, tuple): + encoder_out = encoder_out[0] + + return speech_raw, encoder_out, encoder_out_lens + + def encode2( + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + **kwargs, + ): + """Frontend + Encoder. Note that this method is used by asr_inference.py + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + """ + + ind = kwargs.get("ind", 0) + encoder_out_rm, encoder_out_lens_rm = self.encoder.overlap_chunk_cls.remove_chunk( + encoder_out, + encoder_out_lens, + chunk_outs=None, + ) + # residual_input + encoder_out = torch.cat((speech, encoder_out_rm), dim=-1) + encoder_out_lens = encoder_out_lens_rm + if self.stride_conv is not None: + speech, speech_lengths = self.stride_conv(encoder_out, encoder_out_lens) + if not self.encoder1_encoder2_joint_training: + speech = speech.detach() + speech_lengths = speech_lengths.detach() + # 4. Forward encoder + # feats: (Batch, Length, Dim) + # -> encoder_out: (Batch, Length2, Dim2) + + encoder_out, encoder_out_lens, _ = self.encoder2(speech, speech_lengths, ind=ind) + if isinstance(encoder_out, tuple): + encoder_out = encoder_out[0] + + return encoder_out, encoder_out_lens + + def nll( + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, + ) -> torch.Tensor: + """Compute negative log likelihood(nll) from transformer-decoder + Normally, this function is called in batchify_nll. + Args: + encoder_out: (Batch, Length, Dim) + encoder_out_lens: (Batch,) + ys_pad: (Batch, Length) + ys_pad_lens: (Batch,) + """ + ys_in_pad, ys_out_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id) + ys_in_lens = ys_pad_lens + 1 + + # 1. Forward decoder + decoder_out, _ = self.decoder( + encoder_out, encoder_out_lens, ys_in_pad, ys_in_lens + ) # [batch, seqlen, dim] + batch_size = decoder_out.size(0) + decoder_num_class = decoder_out.size(2) + # nll: negative log-likelihood + nll = torch.nn.functional.cross_entropy( + decoder_out.view(-1, decoder_num_class), + ys_out_pad.view(-1), + ignore_index=self.ignore_id, + reduction="none", + ) + nll = nll.view(batch_size, -1) + nll = nll.sum(dim=1) + assert nll.size(0) == batch_size + return nll + + def batchify_nll( + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, + batch_size: int = 100, + ): + """Compute negative log likelihood(nll) from transformer-decoder + To avoid OOM, this fuction seperate the input into batches. + Then call nll for each batch and combine and return results. + Args: + encoder_out: (Batch, Length, Dim) + encoder_out_lens: (Batch,) + ys_pad: (Batch, Length) + ys_pad_lens: (Batch,) + batch_size: int, samples each batch contain when computing nll, + you may change this to avoid OOM or increase + GPU memory usage + """ + total_num = encoder_out.size(0) + if total_num <= batch_size: + nll = self.nll(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens) + else: + nll = [] + start_idx = 0 + while True: + end_idx = min(start_idx + batch_size, total_num) + batch_encoder_out = encoder_out[start_idx:end_idx, :, :] + batch_encoder_out_lens = encoder_out_lens[start_idx:end_idx] + batch_ys_pad = ys_pad[start_idx:end_idx, :] + batch_ys_pad_lens = ys_pad_lens[start_idx:end_idx] + batch_nll = self.nll( + batch_encoder_out, + batch_encoder_out_lens, + batch_ys_pad, + batch_ys_pad_lens, + ) + nll.append(batch_nll) + start_idx = end_idx + if start_idx == total_num: + break + nll = torch.cat(nll) + assert nll.size(0) == total_num + return nll + + def _calc_att_loss( + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, + ): + ys_in_pad, ys_out_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id) + ys_in_lens = ys_pad_lens + 1 + + # 1. Forward decoder + decoder_out, _ = self.decoder(encoder_out, encoder_out_lens, ys_in_pad, ys_in_lens) + + # 2. Compute attention loss + loss_att = self.criterion_att(decoder_out, ys_out_pad) + acc_att = th_accuracy( + decoder_out.view(-1, self.vocab_size), + ys_out_pad, + ignore_label=self.ignore_id, + ) + + # Compute cer/wer using attention-decoder + if self.training or self.error_calculator is None: + cer_att, wer_att = None, None + else: + ys_hat = decoder_out.argmax(dim=-1) + cer_att, wer_att = self.error_calculator(ys_hat.cpu(), ys_pad.cpu()) + + return loss_att, acc_att, cer_att, wer_att + + def _calc_att_predictor_loss( + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, + ): + ys_in_pad, ys_out_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id) + ys_in_lens = ys_pad_lens + 1 + + encoder_out_mask = sequence_mask( + encoder_out_lens, + maxlen=encoder_out.size(1), + dtype=encoder_out.dtype, + device=encoder_out.device, + )[:, None, :] + mask_chunk_predictor = None + if self.encoder.overlap_chunk_cls is not None: + mask_chunk_predictor = self.encoder.overlap_chunk_cls.get_mask_chunk_predictor( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + mask_shfit_chunk = self.encoder.overlap_chunk_cls.get_mask_shfit_chunk( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + encoder_out = encoder_out * mask_shfit_chunk + pre_acoustic_embeds, pre_token_length, pre_alphas, _ = self.predictor( + encoder_out, + ys_out_pad, + encoder_out_mask, + ignore_id=self.ignore_id, + mask_chunk_predictor=mask_chunk_predictor, + target_label_length=ys_in_lens, + ) + predictor_alignments, predictor_alignments_len = self.predictor.gen_frame_alignments( + pre_alphas, encoder_out_lens + ) + + scama_mask = None + if ( + self.encoder.overlap_chunk_cls is not None + and self.decoder_attention_chunk_type == "chunk" + ): + encoder_chunk_size = self.encoder.overlap_chunk_cls.chunk_size_pad_shift_cur + attention_chunk_center_bias = 0 + attention_chunk_size = encoder_chunk_size + decoder_att_look_back_factor = ( + self.encoder.overlap_chunk_cls.decoder_att_look_back_factor_cur + ) + mask_shift_att_chunk_decoder = ( + self.encoder.overlap_chunk_cls.get_mask_shift_att_chunk_decoder( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + ) + scama_mask = self.build_scama_mask_for_cross_attention_decoder_fn( + predictor_alignments=predictor_alignments, + encoder_sequence_length=encoder_out_lens, + chunk_size=1, + encoder_chunk_size=encoder_chunk_size, + attention_chunk_center_bias=attention_chunk_center_bias, + attention_chunk_size=attention_chunk_size, + attention_chunk_type=self.decoder_attention_chunk_type, + step=None, + predictor_mask_chunk_hopping=mask_chunk_predictor, + decoder_att_look_back_factor=decoder_att_look_back_factor, + mask_shift_att_chunk_decoder=mask_shift_att_chunk_decoder, + target_length=ys_in_lens, + is_training=self.training, + ) + elif self.encoder.overlap_chunk_cls is not None: + encoder_out, encoder_out_lens = self.encoder.overlap_chunk_cls.remove_chunk( + encoder_out, encoder_out_lens, chunk_outs=None + ) + # try: + # 1. Forward decoder + decoder_out, _ = self.decoder( + encoder_out, + encoder_out_lens, + ys_in_pad, + ys_in_lens, + chunk_mask=scama_mask, + pre_acoustic_embeds=pre_acoustic_embeds, + ) + + # 2. Compute attention loss + loss_att = self.criterion_att(decoder_out, ys_out_pad) + acc_att = th_accuracy( + decoder_out.view(-1, self.vocab_size), + ys_out_pad, + ignore_label=self.ignore_id, + ) + # predictor loss + loss_pre = self.criterion_pre(ys_in_lens.type_as(pre_token_length), pre_token_length) + # Compute cer/wer using attention-decoder + if self.training or self.error_calculator is None: + cer_att, wer_att = None, None + else: + ys_hat = decoder_out.argmax(dim=-1) + cer_att, wer_att = self.error_calculator(ys_hat.cpu(), ys_pad.cpu()) + + return loss_att, acc_att, cer_att, wer_att, loss_pre + + def _calc_att_predictor_loss2( + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, + ): + ys_in_pad, ys_out_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id) + ys_in_lens = ys_pad_lens + 1 + + encoder_out_mask = sequence_mask( + encoder_out_lens, + maxlen=encoder_out.size(1), + dtype=encoder_out.dtype, + device=encoder_out.device, + )[:, None, :] + mask_chunk_predictor = None + if self.encoder2.overlap_chunk_cls is not None: + mask_chunk_predictor = self.encoder2.overlap_chunk_cls.get_mask_chunk_predictor( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + mask_shfit_chunk = self.encoder2.overlap_chunk_cls.get_mask_shfit_chunk( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + encoder_out = encoder_out * mask_shfit_chunk + pre_acoustic_embeds, pre_token_length, pre_alphas, _ = self.predictor2( + encoder_out, + ys_out_pad, + encoder_out_mask, + ignore_id=self.ignore_id, + mask_chunk_predictor=mask_chunk_predictor, + target_label_length=ys_in_lens, + ) + predictor_alignments, predictor_alignments_len = self.predictor2.gen_frame_alignments( + pre_alphas, encoder_out_lens + ) + + scama_mask = None + if ( + self.encoder2.overlap_chunk_cls is not None + and self.decoder_attention_chunk_type2 == "chunk" + ): + encoder_chunk_size = self.encoder2.overlap_chunk_cls.chunk_size_pad_shift_cur + attention_chunk_center_bias = 0 + attention_chunk_size = encoder_chunk_size + decoder_att_look_back_factor = ( + self.encoder2.overlap_chunk_cls.decoder_att_look_back_factor_cur + ) + mask_shift_att_chunk_decoder = ( + self.encoder2.overlap_chunk_cls.get_mask_shift_att_chunk_decoder( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + ) + scama_mask = self.build_scama_mask_for_cross_attention_decoder_fn2( + predictor_alignments=predictor_alignments, + encoder_sequence_length=encoder_out_lens, + chunk_size=1, + encoder_chunk_size=encoder_chunk_size, + attention_chunk_center_bias=attention_chunk_center_bias, + attention_chunk_size=attention_chunk_size, + attention_chunk_type=self.decoder_attention_chunk_type2, + step=None, + predictor_mask_chunk_hopping=mask_chunk_predictor, + decoder_att_look_back_factor=decoder_att_look_back_factor, + mask_shift_att_chunk_decoder=mask_shift_att_chunk_decoder, + target_length=ys_in_lens, + is_training=self.training, + ) + elif self.encoder2.overlap_chunk_cls is not None: + encoder_out, encoder_out_lens = self.encoder2.overlap_chunk_cls.remove_chunk( + encoder_out, encoder_out_lens, chunk_outs=None + ) + # try: + # 1. Forward decoder + decoder_out, _ = self.decoder2( + encoder_out, + encoder_out_lens, + ys_in_pad, + ys_in_lens, + chunk_mask=scama_mask, + pre_acoustic_embeds=pre_acoustic_embeds, + ) + + # 2. Compute attention loss + loss_att = self.criterion_att(decoder_out, ys_out_pad) + acc_att = th_accuracy( + decoder_out.view(-1, self.vocab_size), + ys_out_pad, + ignore_label=self.ignore_id, + ) + # predictor loss + loss_pre = self.criterion_pre(ys_in_lens.type_as(pre_token_length), pre_token_length) + # Compute cer/wer using attention-decoder + if self.training or self.error_calculator is None: + cer_att, wer_att = None, None + else: + ys_hat = decoder_out.argmax(dim=-1) + cer_att, wer_att = self.error_calculator(ys_hat.cpu(), ys_pad.cpu()) + + return loss_att, acc_att, cer_att, wer_att, loss_pre + + def calc_predictor_mask( + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor = None, + ys_pad_lens: torch.Tensor = None, + ): + # ys_in_pad, ys_out_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id) + # ys_in_lens = ys_pad_lens + 1 + ys_out_pad, ys_in_lens = None, None + + encoder_out_mask = sequence_mask( + encoder_out_lens, + maxlen=encoder_out.size(1), + dtype=encoder_out.dtype, + device=encoder_out.device, + )[:, None, :] + mask_chunk_predictor = None + if self.encoder.overlap_chunk_cls is not None: + mask_chunk_predictor = self.encoder.overlap_chunk_cls.get_mask_chunk_predictor( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + mask_shfit_chunk = self.encoder.overlap_chunk_cls.get_mask_shfit_chunk( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + encoder_out = encoder_out * mask_shfit_chunk + pre_acoustic_embeds, pre_token_length, pre_alphas, _ = self.predictor( + encoder_out, + ys_out_pad, + encoder_out_mask, + ignore_id=self.ignore_id, + mask_chunk_predictor=mask_chunk_predictor, + target_label_length=ys_in_lens, + ) + predictor_alignments, predictor_alignments_len = self.predictor.gen_frame_alignments( + pre_alphas, encoder_out_lens + ) + + scama_mask = None + if ( + self.encoder.overlap_chunk_cls is not None + and self.decoder_attention_chunk_type == "chunk" + ): + encoder_chunk_size = self.encoder.overlap_chunk_cls.chunk_size_pad_shift_cur + attention_chunk_center_bias = 0 + attention_chunk_size = encoder_chunk_size + decoder_att_look_back_factor = ( + self.encoder.overlap_chunk_cls.decoder_att_look_back_factor_cur + ) + mask_shift_att_chunk_decoder = ( + self.encoder.overlap_chunk_cls.get_mask_shift_att_chunk_decoder( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + ) + scama_mask = self.build_scama_mask_for_cross_attention_decoder_fn( + predictor_alignments=predictor_alignments, + encoder_sequence_length=encoder_out_lens, + chunk_size=1, + encoder_chunk_size=encoder_chunk_size, + attention_chunk_center_bias=attention_chunk_center_bias, + attention_chunk_size=attention_chunk_size, + attention_chunk_type=self.decoder_attention_chunk_type, + step=None, + predictor_mask_chunk_hopping=mask_chunk_predictor, + decoder_att_look_back_factor=decoder_att_look_back_factor, + mask_shift_att_chunk_decoder=mask_shift_att_chunk_decoder, + target_length=ys_in_lens, + is_training=self.training, + ) + elif self.encoder.overlap_chunk_cls is not None: + encoder_out, encoder_out_lens = self.encoder.overlap_chunk_cls.remove_chunk( + encoder_out, encoder_out_lens, chunk_outs=None + ) + + return ( + pre_acoustic_embeds, + pre_token_length, + predictor_alignments, + predictor_alignments_len, + scama_mask, + ) + + def calc_predictor_mask2( + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor = None, + ys_pad_lens: torch.Tensor = None, + ): + # ys_in_pad, ys_out_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id) + # ys_in_lens = ys_pad_lens + 1 + ys_out_pad, ys_in_lens = None, None + + encoder_out_mask = sequence_mask( + encoder_out_lens, + maxlen=encoder_out.size(1), + dtype=encoder_out.dtype, + device=encoder_out.device, + )[:, None, :] + mask_chunk_predictor = None + if self.encoder2.overlap_chunk_cls is not None: + mask_chunk_predictor = self.encoder2.overlap_chunk_cls.get_mask_chunk_predictor( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + mask_shfit_chunk = self.encoder2.overlap_chunk_cls.get_mask_shfit_chunk( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + encoder_out = encoder_out * mask_shfit_chunk + + with torch.cuda.amp.autocast(enabled=False): + pre_acoustic_embeds, pre_token_length, pre_alphas, _ = self.predictor2( + encoder_out, + ys_out_pad, + encoder_out_mask, + ignore_id=self.ignore_id, + mask_chunk_predictor=mask_chunk_predictor, + target_label_length=ys_in_lens, + ) + predictor_alignments, predictor_alignments_len = self.predictor2.gen_frame_alignments( + pre_alphas, encoder_out_lens + ) + + scama_mask = None + if ( + self.encoder2.overlap_chunk_cls is not None + and self.decoder_attention_chunk_type2 == "chunk" + ): + encoder_chunk_size = self.encoder2.overlap_chunk_cls.chunk_size_pad_shift_cur + attention_chunk_center_bias = 0 + attention_chunk_size = encoder_chunk_size + decoder_att_look_back_factor = ( + self.encoder2.overlap_chunk_cls.decoder_att_look_back_factor_cur + ) + mask_shift_att_chunk_decoder = ( + self.encoder2.overlap_chunk_cls.get_mask_shift_att_chunk_decoder( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + ) + scama_mask = self.build_scama_mask_for_cross_attention_decoder_fn2( + predictor_alignments=predictor_alignments, + encoder_sequence_length=encoder_out_lens, + chunk_size=1, + encoder_chunk_size=encoder_chunk_size, + attention_chunk_center_bias=attention_chunk_center_bias, + attention_chunk_size=attention_chunk_size, + attention_chunk_type=self.decoder_attention_chunk_type2, + step=None, + predictor_mask_chunk_hopping=mask_chunk_predictor, + decoder_att_look_back_factor=decoder_att_look_back_factor, + mask_shift_att_chunk_decoder=mask_shift_att_chunk_decoder, + target_length=ys_in_lens, + is_training=self.training, + ) + elif self.encoder2.overlap_chunk_cls is not None: + encoder_out, encoder_out_lens = self.encoder2.overlap_chunk_cls.remove_chunk( + encoder_out, encoder_out_lens, chunk_outs=None + ) + + return ( + pre_acoustic_embeds, + pre_token_length, + predictor_alignments, + predictor_alignments_len, + scama_mask, + ) + + def init_beam_search( + self, + **kwargs, + ): + from funasr.models.uniasr.beam_search import BeamSearchScama + from funasr.models.transformer.scorers.ctc import CTCPrefixScorer + from funasr.models.transformer.scorers.length_bonus import LengthBonus + + decoding_mode = kwargs.get("decoding_mode", "model1") + if decoding_mode == "model1": + decoder = self.decoder + else: + decoder = self.decoder2 + # 1. Build ASR model + scorers = {} + + if self.ctc != None: + ctc = CTCPrefixScorer(ctc=self.ctc, eos=self.eos) + scorers.update(ctc=ctc) + token_list = kwargs.get("token_list") + scorers.update( + decoder=decoder, + length_bonus=LengthBonus(len(token_list)), + ) + + # 3. Build ngram model + # ngram is not supported now + ngram = None + scorers["ngram"] = ngram + + weights = dict( + decoder=1.0 - kwargs.get("decoding_ctc_weight", 0.0), + ctc=kwargs.get("decoding_ctc_weight", 0.0), + lm=kwargs.get("lm_weight", 0.0), + ngram=kwargs.get("ngram_weight", 0.0), + length_bonus=kwargs.get("penalty", 0.0), + ) + beam_search = BeamSearchScama( + beam_size=kwargs.get("beam_size", 5), + weights=weights, + scorers=scorers, + sos=self.sos, + eos=self.eos, + vocab_size=len(token_list), + token_list=token_list, + pre_beam_score_key=None if self.ctc_weight == 1.0 else "full", + ) + + self.beam_search = beam_search + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): + + decoding_model = kwargs.get("decoding_model", "normal") + token_num_relax = kwargs.get("token_num_relax", 5) + if decoding_model == "fast": + decoding_ind = 0 + decoding_mode = "model1" + elif decoding_model == "offline": + decoding_ind = 1 + decoding_mode = "model2" + else: + decoding_ind = 0 + decoding_mode = "model2" + # init beamsearch + + if self.beam_search is None: + logging.info("enable beam_search") + self.init_beam_search(decoding_mode=decoding_mode, **kwargs) + self.nbest = kwargs.get("nbest", 1) + + meta_data = {} + if ( + isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank" + ): # fbank + speech, speech_lengths = data_in, data_lengths + if len(speech.shape) < 3: + speech = speech[None, :, :] + if speech_lengths is 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), + data_type=kwargs.get("data_type", "sound"), + tokenizer=tokenizer, + ) + time2 = time.perf_counter() + meta_data["load_data"] = f"{time2 - time1:0.3f}" + speech, speech_lengths = extract_fbank( + audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend + ) + time3 = time.perf_counter() + meta_data["extract_feat"] = f"{time3 - time2:0.3f}" + meta_data["batch_data_time"] = ( + speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 + ) + + speech = speech.to(device=kwargs["device"]) + speech_lengths = speech_lengths.to(device=kwargs["device"]) + speech_raw = speech.clone().to(device=kwargs["device"]) + # Encoder + _, encoder_out, encoder_out_lens = self.encode(speech, speech_lengths, ind=decoding_ind) + if decoding_mode == "model1": + predictor_outs = self.calc_predictor_mask(encoder_out, encoder_out_lens) + else: + encoder_out, encoder_out_lens = self.encode2( + encoder_out, encoder_out_lens, speech_raw, speech_lengths, ind=decoding_ind + ) + predictor_outs = self.calc_predictor_mask2(encoder_out, encoder_out_lens) + + scama_mask = predictor_outs[4] + pre_token_length = predictor_outs[1] + pre_acoustic_embeds = predictor_outs[0] + maxlen = pre_token_length.sum().item() + token_num_relax + minlen = max(0, pre_token_length.sum().item() - token_num_relax) + # c. Passed the encoder result and the beam search + nbest_hyps = self.beam_search( + x=encoder_out[0], + scama_mask=scama_mask, + pre_acoustic_embeds=pre_acoustic_embeds, + maxlenratio=0.0, + minlenratio=0.0, + maxlen=int(maxlen), + minlen=int(minlen), + ) + + nbest_hyps = nbest_hyps[: self.nbest] + + results = [] + for hyp in nbest_hyps: + + # 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 != 0, token_int)) + + # Change integer-ids to tokens + token = tokenizer.ids2tokens(token_int) + text_postprocessed = tokenizer.tokens2text(token) + if not hasattr(tokenizer, "bpemodel"): + text_postprocessed, _ = postprocess_utils.sentence_postprocess(token) + + result_i = {"key": key[0], "text": text_postprocessed} + results.append(result_i) + + return results, meta_data diff --git a/replaced_files/funasr_nano_model.py b/replaced_files/funasr_nano_model.py new file mode 100644 index 0000000..73da9f1 --- /dev/null +++ b/replaced_files/funasr_nano_model.py @@ -0,0 +1,746 @@ +import logging +import os +import random +import re +import string +import time +import traceback +from typing import Union + +import torch +import torch.nn as nn + +from funasr.metrics.compute_acc import compute_accuracy +from funasr.register import tables +from funasr.train_utils.device_funcs import force_gatherable, to_device +from funasr.utils.datadir_writer import DatadirWriter +from funasr.utils.load_utils import extract_fbank, load_audio_text_image_video +from transformers import AutoConfig, AutoModelForCausalLM + +from funasr.models.fun_asr_nano.ctc import CTC +from funasr.models.fun_asr_nano.tools.utils import forced_align + +dtype_map = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32} + + +@tables.register("model_classes", "FunASRNano") +class FunASRNano(nn.Module): + def __init__( + self, + audio_encoder: str = None, + audio_encoder_conf: dict = None, + audio_adaptor: str = None, + audio_adaptor_conf: dict = None, + llm: str = None, + llm_conf: dict = None, + input_size: int = 80, + length_normalized_loss: bool = False, + **kwargs, + ): + super().__init__() + + # audio encoder + hub = audio_encoder_conf.get("hub", None) + self.audio_encoder_activation_checkpoint = audio_encoder_conf.get( + "activation_checkpoint", False + ) + if hub == "ms": + from funasr import AutoModel + + model = AutoModel(model=audio_encoder, model_revision="master") + audio_encoder_output_size = ( + model.model.encoder_output_size + if hasattr(model.model, "encoder_output_size") + else -1 + ) + audio_encoder = ( + model.model.model.encoder if hasattr(model.model, "model") else model.model.encoder + ) + else: + encoder_class = tables.encoder_classes.get(audio_encoder) + audio_encoder = encoder_class(input_size=input_size, **audio_encoder_conf) + audio_encoder_output_size = audio_encoder.output_size() + freeze = audio_encoder_conf.get("freeze", True) + + if freeze: + for _, param in audio_encoder.named_parameters(): + param.requires_grad = False + audio_encoder.eval() + self.audio_encoder = audio_encoder + + # llm + self.llm = None + init_param_path = llm_conf.get("init_param_path", None) + llm_dim = None + + llm_load_kwargs = llm_conf.get("load_kwargs", {}) + config = AutoConfig.from_pretrained(init_param_path) + model = AutoModelForCausalLM.from_config(config, **llm_load_kwargs) + + freeze = llm_conf.get("freeze", True) + if freeze: + for _, param in model.named_parameters(): + param.requires_grad = False + model.eval() + if llm_conf.get("activation_checkpoint", False): + model.gradient_checkpointing_enable() + + self.llm_dtype = llm_conf.get("llm_dtype", "fp32") + self.llm = model.to(dtype_map[self.llm_dtype]) + llm_dim = model.get_input_embeddings().weight.shape[-1] + + # adaptor + adaptor_class = tables.adaptor_classes.get(audio_adaptor) + if audio_encoder_output_size > 0: + audio_adaptor_conf["encoder_dim"] = audio_encoder_output_size + audio_adaptor_conf["llm_dim"] = ( + llm_dim if llm_dim is not None else audio_adaptor_conf["llm_dim"] + ) + audio_adaptor = adaptor_class(**audio_adaptor_conf) + freeze = audio_adaptor_conf.get("freeze", False) + if freeze: + for _, param in audio_adaptor.named_parameters(): + param.requires_grad = False + audio_adaptor.eval() + self.audio_adaptor = audio_adaptor + self.use_low_frame_rate = audio_adaptor_conf.get("use_low_frame_rate", False) + + # ctc decoder + self.ctc_decoder = None + # TODO: fix table name + ctc_decoder_class = tables.adaptor_classes.get(kwargs.get("ctc_decoder", None)) + if ctc_decoder_class is not None: + ctc_tokenizer = ( + kwargs.get("ctc_tokenizer", None) + if "ctc_tokenizer" in kwargs + else kwargs["dataset_conf"]["ctc_tokenizer"] + ) + ctc_tokenizer_conf = ( + kwargs.get("ctc_tokenizer_conf", None) + if "ctc_tokenizer_conf" in kwargs + else kwargs["dataset_conf"]["ctc_tokenizer_conf"] + ) + if ctc_tokenizer is not None and ctc_tokenizer_conf is not None: + ctc_tokenizer_class = tables.tokenizer_classes.get(ctc_tokenizer) + ctc_tokenizer = ctc_tokenizer_class(**ctc_tokenizer_conf) + self.ctc_tokenizer = ctc_tokenizer + assert ctc_tokenizer is not None, f"ctc_tokenizer must be set" + + ctc_vocab_size = kwargs.get("ctc_vocab_size", 60515) + ctc_decoder_conf = kwargs.get("ctc_decoder_conf", {}) + if audio_encoder_output_size > 0: + ctc_decoder_conf["encoder_dim"] = audio_encoder_output_size + self.ctc_decoder = ctc_decoder_class(**ctc_decoder_conf) + init_param_path = ctc_decoder_conf.get("init_param_path", None) + if init_param_path is not None: + src_state = torch.load(init_param_path, map_location="cpu") + flag = self.ctc_decoder.load_state_dict(src_state, strict=False) + logging.info(f"Loading ctc_decoder ckpt: {init_param_path}, status: {flag}") + freeze = ctc_decoder_conf.get("freeze", False) + if freeze: + for _, param in self.ctc_decoder.named_parameters(): + param.requires_grad = False + self.ctc_decoder.eval() + + ctc_conf = kwargs.get("ctc_conf", {}) + self.blank_id = ctc_conf.get("blank_id", ctc_vocab_size - 1) + self.ctc_weight = kwargs.get("ctc_weight", 0.3) + self.ctc = CTC( + odim=ctc_vocab_size, + encoder_output_size=audio_encoder_output_size, + blank_id=self.blank_id, + **ctc_conf, + ) + self.detach_ctc_decoder = kwargs.get("detach_ctc_decoder", True) + self.error_calculator = None + + self.length_normalized_loss = length_normalized_loss + rank = int(os.environ.get("RANK", 0)) + logging.info(f"rank: {rank}, model is builded.") + + def forward( + self, + speech: torch.Tensor = None, + speech_lengths: torch.Tensor = None, + input_ids: torch.Tensor = None, + attention_mask: torch.Tensor = None, + labels_ids: torch.Tensor = None, + fbank_beg: torch.Tensor = None, + fbank_mask: torch.Tensor = None, + **kwargs, + ): + batch_size, token_num = input_ids.shape + stats = {} + input_ids[input_ids < 0] = 0 + inputs_embeds = self.llm.model.get_input_embeddings()(input_ids) + if speech is not None: + if len(speech_lengths.size()) > 1: + speech_lengths = speech_lengths[:, 0] + batch_size_speech, frames, _ = speech.shape + + # audio encoder + if self.audio_encoder_activation_checkpoint: + from torch.utils.checkpoint import checkpoint + + encoder_out, encoder_out_lens = checkpoint( + self.encode, speech, speech_lengths, use_reentrant=False + ) + else: + encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) + + # audio_adaptor + encoder_out, encoder_out_lens = self.audio_adaptor(encoder_out, encoder_out_lens) + + batch_size, token_num, dims = inputs_embeds.shape + fake_token_len = kwargs.get("fake_token_len") + fake_token_len[fake_token_len < 0] = 0 + fbank_beg[fbank_beg < 0] = 0 + + speech_idx = 0 + for batch_idx in range(batch_size): + for turn_id in range(fbank_beg.shape[1]): + fbank_beg_idx = fbank_beg[batch_idx, turn_id].item() + if fbank_beg_idx > 0: + speech_token_len = fake_token_len[batch_idx, turn_id] + speech_token = encoder_out[speech_idx, :speech_token_len, :] + + try: + inputs_embeds[ + batch_idx, + fbank_beg_idx : fbank_beg_idx + speech_token_len, + :, + ] = speech_token + except Exception as e: + logging.error(f"{str(e)}, {traceback.format_exc()}") + logging.info( + f"batch_idx: {batch_idx}, inputs_embeds: {inputs_embeds.shape}, fbank_beg_idx: {fbank_beg_idx}, speech_token_len: {speech_token_len}, encoder_out: {encoder_out.shape}, encoder_out_lens: {encoder_out_lens}, fake_token_len: {fake_token_len}, speech_lengths: {speech_lengths}" + ) + speech_token_len = encoder_out_lens[speech_idx].item() + speech_token = encoder_out[speech_idx, :speech_token_len, :] + inputs_embeds[ + batch_idx, + fbank_beg_idx : fbank_beg_idx + speech_token_len, + :, + ] = speech_token + + speech_idx += 1 + + stats["batch_size_speech"] = batch_size_speech + stats["batch_size_x_frames"] = frames * batch_size_speech + stats["batch_size_real_frames"] = speech_lengths.sum().item() + stats["padding_frames"] = stats["batch_size_x_frames"] - stats["batch_size_real_frames"] + + device_type = next(self.parameters()).device.type + with torch.autocast( + device_type=device_type if device_type in ["cuda", "xpu", "mps"] else "cpu", + enabled=True if self.llm_dtype != "fp32" else False, + dtype=dtype_map[self.llm_dtype], + ): + labels_ids[labels_ids == -1] = -100 + attention_mask[attention_mask < 0] = 0 + model_outputs = self.llm( + inputs_embeds=inputs_embeds.to(dtype_map[self.llm_dtype]), + attention_mask=attention_mask, + labels=labels_ids, + ) + loss = model_outputs.loss + + with torch.no_grad(): + preds = torch.argmax(model_outputs.logits, -1) + acc_att = compute_accuracy(preds[:, :-1], labels_ids[:, 1:], ignore_label=-100) + stats["acc"] = acc_att + + stats["loss"] = torch.clone(loss.detach()) + stats["batch_size"] = batch_size + + stats["batch_size_x_tokens"] = token_num * batch_size + stats["batch_size_real_tokens"] = attention_mask.sum().item() + stats["padding_tokens"] = stats["batch_size_x_tokens"] - stats["batch_size_real_tokens"] + + dialog_turns = (fbank_beg > 0).sum(-1) + dialog_turns_max = torch.max(dialog_turns).int().item() + dialog_turns_avg = dialog_turns.sum().item() / batch_size + stats["dialog_turns_max"] = dialog_turns_max + stats["dialog_turns_avg"] = dialog_turns_avg + + # force_gatherable: to-device and to-tensor if scalar for DataParallel + if self.length_normalized_loss: + batch_size = int((labels_ids > 0 + 1).sum()) + loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) + return loss, stats, weight + + def forward_export(self, speech, speech_lengths, **kwargs): + x, olens = self.audio_encoder(speech, speech_lengths) + encoder_out, encoder_out_lens = self.audio_adaptor(x, olens) + return encoder_out, encoder_out_lens + + def encode(self, speech, speech_lengths): + # audio encoder + encoder_out, encoder_out_lens = self.audio_encoder(speech, speech_lengths) + + return encoder_out, encoder_out_lens + + def data_template(self, data): + system, user, assistant = [], [], [] + for i, item in enumerate(data): + role = item["role"] + content = item["content"] + if role == "system": + system.append(content) + elif role == "user": + if "audio" in item: + audio = item["audio"] + content = [content, audio] + user.append(content) + elif role == "assistant": + assistant.append(content) + + system = system * len(user) + + contents = { + "system": system, + "user": user, + "assistant": assistant, + } + + return contents + + def data_load_speech(self, contents: dict, tokenizer, frontend, meta_data={}, **kwargs): + system = contents["system"] + user = contents["user"] + assistant = contents["assistant"] + pattern = re.compile(r"(<\|startofspeech\|>.*?<\|endofspeech\|>)") + do_think = True + sys_prompt = True + if "dataset_conf" in kwargs: + do_think = kwargs["dataset_conf"].get("do_think", True) + sys_prompt = kwargs["dataset_conf"].get("sys_prompt", True) + + input_ids, labels, fbank, fbank_lens, fbank_mask, fbank_beg, fake_token_len = ( + [], + [], + [], + [], + [], + [], + [], + ) + input_source_ids = [] + for i, (system_prompt, user_prompt, target_out) in enumerate(zip(system, user, assistant)): + if i >= kwargs.get("multiturn_num_max", 5): + break + if len(input_ids) > kwargs.get("max_token_length", 1500): + break + if isinstance(user_prompt, (list, tuple)): + user_prompt, audio = user_prompt + if i == 0: + if kwargs.get("infer_with_assistant_input", False): + source_input = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{user_prompt}" + if not sys_prompt: + source_input = f"<|im_start|>user\n{user_prompt}" + else: + 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" + if not sys_prompt: + source_input = ( + f"<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n" + ) + else: + if kwargs.get("infer_with_assistant_input", False): + source_input = f"<|im_start|>user\n{user_prompt}" + else: + source_input = ( + f"<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n" + ) + if not do_think: + source_input += "\n\n\n\n" + if kwargs.get("prev_text", None) is not None: + source_input += kwargs["prev_text"] + + splits = pattern.split(source_input) + source_ids = [] + fbank_mask_i = [] + fake_token_len_i = 0 + fbank_beg_i = -1 + speech, speech_lengths = [], [] + for k, sub_str in enumerate(splits): + if not sub_str.startswith("<|startofspeech|>"): + sub_token = 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("!"): + sub_str = sub_str[1:] + if sub_str.startswith("!"): # !!: audio sample point + sub_str = audio + try: + time1 = time.perf_counter() + data_src = load_audio_text_image_video( + sub_str, fs=frontend.fs, **kwargs + ) + time2 = time.perf_counter() + meta_data["load_data"] = f"{time2 - time1:0.3f}" + except Exception as e: + logging.error(f"Loading wav failed! {str(e)}, {traceback.format_exc()}") + + speech, speech_lengths = extract_fbank( + data_src, + data_type=kwargs.get("data_type", "sound"), + frontend=frontend, + is_final=True, + ) # speech: [b, T, d] + + 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 + ) + + if self.use_low_frame_rate: + 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 + fbank_mask_i += [1] * len(fake_token) + + 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 = tokenizer.encode(target_out) + input_source_ids = input_ids + source_ids + input_ids += source_ids + target_ids + labels += source_mask + target_ids + fbank_mask += fbank_mask_i + if len(speech) > 0: + fbank.append(speech[0, :, :]) + fbank_lens.append(speech_lengths) + + 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_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) + source_ids = torch.tensor(input_source_ids, dtype=torch.int64) + target_ids = torch.tensor(target_ids, dtype=torch.int64) + + if len(fbank) > 0: + speech = torch.nn.utils.rnn.pad_sequence(fbank, batch_first=True, padding_value=0.0) + speech_lengths = torch.nn.utils.rnn.pad_sequence( + fbank_lens, batch_first=True, padding_value=-1 + ) + else: + speech = [] + speech_lengths = [] + output = { + "speech": speech, + "speech_lengths": speech_lengths, + "fbank_mask": fbank_mask[None, :], + "fbank_beg": fbank_beg[None,], + "fake_token_len": fake_token_len[None, :], + "input_ids": input_ids[None,], + "attention_mask": attention_mask[None,], + "labels_ids": labels, + "source_ids": source_ids[None, :], + "target_ids": target_ids[None, :], + } + + return output + + def inference_prepare( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): + meta_data = {} + + if kwargs.get("batch_size", 1) > 1: + raise NotImplementedError("batch decoding is not implemented") + + contents = self.data_template(data_in[0]) + output = self.data_load_speech(contents, tokenizer, frontend, meta_data=meta_data, **kwargs) + batch = to_device(output, kwargs["device"]) + + # audio encoder + speech = batch["speech"] + + if len(speech) > 0: + if "audio_embedding" in kwargs and "audio_embedding_lens" in kwargs: + encoder_out = kwargs["audio_embedding"] + encoder_out_lens = kwargs["audio_embedding_lens"] + else: + speech_lengths = batch["speech_lengths"][:, 0] + # fp16 + if kwargs.get("fp16", False): + speech = speech.to(torch.float16) + elif kwargs.get("bf16", False): + speech = speech.to(torch.bfloat16) + # audio encoder + encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) + + # audio_adaptor + adaptor_out, adaptor_out_lens = self.audio_adaptor(encoder_out, encoder_out_lens) + meta_data["encoder_out"] = encoder_out + meta_data["encoder_out_lens"] = encoder_out_lens + meta_data["audio_adaptor_out"] = adaptor_out + meta_data["audio_adaptor_out_lens"] = adaptor_out_lens + + input_ids = batch["input_ids"] + source_ids = batch["source_ids"] + fbank_beg = batch["fbank_beg"] + fake_token_len = batch["fake_token_len"] + + if not kwargs.get("teacherforcing", False): + input_ids = source_ids + + input_ids[input_ids < 0] = 0 + inputs_embeds = self.llm.model.get_input_embeddings()(input_ids) + + batch_size, token_num, dims = inputs_embeds.shape + + fake_token_len[fake_token_len < 0] = 0 + fbank_beg[fbank_beg < 0] = 0 + + speech_idx = 0 + for batch_idx in range(batch_size): + for turn_id in range(fbank_beg.shape[1]): + fbank_beg_idx = fbank_beg[batch_idx, turn_id].item() + if fbank_beg_idx > 0: + speech_token_len = fake_token_len[batch_idx, turn_id] + speech_token = adaptor_out[speech_idx, :speech_token_len, :] + + try: + inputs_embeds[ + batch_idx, + fbank_beg_idx : fbank_beg_idx + speech_token_len, + :, + ] = speech_token + except Exception as e: + # + logging.error(f"{str(e)}, {traceback.format_exc()}") + logging.info( + f"batch_idx: {batch_idx}, inputs_embeds: {inputs_embeds.shape}, fbank_beg_idx: {fbank_beg_idx}, speech_token_len: {speech_token_len}, adaptor_out: {adaptor_out.shape}, adaptor_out_lens: {adaptor_out_lens}, fake_token_len: {fake_token_len}, speech_lengths: {speech_lengths}" + ) + speech_token_len = adaptor_out_lens[speech_idx].item() + speech_token = adaptor_out[speech_idx, :speech_token_len, :] + inputs_embeds[ + batch_idx, + fbank_beg_idx : fbank_beg_idx + speech_token_len, + :, + ] = speech_token + + speech_idx += 1 + return inputs_embeds, contents, batch, source_ids, meta_data + + def get_prompt(self, hotwords: list[str], language: str = None, itn: bool = True): + if len(hotwords) > 0: + hotwords = ", ".join(hotwords) + prompt = f"请结合上下文信息，更加准确地完成语音转写任务。如果没有相关信息，我们会留空。\n\n\n**上下文信息：**\n\n\n" + prompt += f"热词列表：[{hotwords}]\n" + else: + prompt = "" + if language is None: + prompt += "语音转写" + else: + prompt += f"语音转写成{language}" + if not itn: + prompt += "，不进行文本规整" + return prompt + "：" + + def generate_chatml(self, prompt: str, data: Union[str, torch.Tensor]): + if isinstance(data, str): + return [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": f"{prompt}<|startofspeech|>!{data}<|endofspeech|>"}, + {"role": "assistant", "content": "null"}, + ] + elif isinstance(data, torch.Tensor): + return [ + {"role": "system", "content": "You are a helpful assistant."}, + { + "role": "user", + "content": f"{prompt}<|startofspeech|>!!<|endofspeech|>", + "audio": data, + }, + {"role": "assistant", "content": "null"}, + ] + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): + prompt = self.get_prompt( + kwargs.get("hotwords", []), kwargs.get("language", None), kwargs.get("itn", True) + ) + data_in = [self.generate_chatml(prompt, data) for data in data_in] + + if key is None: + key = [] + for _ in data_in: + chars = string.ascii_letters + string.digits + key.append("rand_key_" + "".join(random.choice(chars) for _ in range(13))) + + return self.inference_llm( + data_in, + data_lengths=data_lengths, + key=key, + tokenizer=tokenizer, + frontend=frontend, + **kwargs, + ) + + def inference_llm( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): + inputs_embeds, contents, batch, source_ids, meta_data = self.inference_prepare( + data_in, data_lengths, key, tokenizer, frontend, **kwargs + ) + + ctc_results = [] + if self.ctc_decoder is not None: + encoder_out = meta_data["encoder_out"] + encoder_out_lens = meta_data["encoder_out_lens"] + decoder_out, decoder_out_lens = self.ctc_decoder(encoder_out, encoder_out_lens) + ctc_logits = self.ctc.log_softmax(decoder_out) + + b, n, d = encoder_out.size() + if isinstance(key[0], (list, tuple)): + key = key[0] + if len(key) < b: + key = key * b + for i in range(b): + x = ctc_logits[i, : encoder_out_lens[i].item(), :] + yseq = x.argmax(dim=-1) + yseq = torch.unique_consecutive(yseq, dim=-1) + mask = yseq != self.blank_id + token_int = yseq[mask].tolist() + # Change integer-ids to tokens + text = self.ctc_tokenizer.decode(token_int) + ctc_results.append({"key": key[i], "text": text, "ctc_logits": x}) + + llm_dtype = kwargs.get("llm_dtype", "fp32") + if llm_dtype == "fp32": + llm_dtype = "fp16" if kwargs.get("fp16", False) else llm_dtype + llm_dtype = "bf16" if kwargs.get("bf16", False) else llm_dtype + + device_type = torch.device(kwargs.get("device", "cuda")).type + with torch.autocast( + device_type=device_type if device_type in ["cuda", "xpu", "mps"] else "cpu", + enabled=True if llm_dtype != "fp32" else False, + dtype=dtype_map[llm_dtype], + ): + label = contents["assistant"][-1] + self.llm = self.llm.to(dtype_map[llm_dtype]) + inputs_embeds = inputs_embeds.to(dtype_map[llm_dtype]) + llm_kwargs = kwargs.get("llm_kwargs", {}) + if not kwargs.get("teacherforcing", False): + attention_mask = batch.get("attention_mask", None) + generated_ids = self.llm.generate( + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + max_new_tokens=kwargs.get("max_length", 512), + pad_token_id=self.llm.config.pad_token_id or self.llm.config.eos_token_id, + **llm_kwargs, + ) + + response = tokenizer.batch_decode( + generated_ids, + skip_special_tokens=kwargs.get("skip_special_tokens", True), + )[0] + + loss = None + else: + labels_ids = batch["labels_ids"] + labels_ids[labels_ids == -1] = -100 + attention_mask = batch.get("attention_mask", None) + model_outputs = self.llm( + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + labels=labels_ids, + pad_token_id=self.llm.config.pad_token_id or self.llm.config.eos_token_id, + **llm_kwargs, + ) + + preds = torch.argmax(model_outputs.logits, -1)[:, source_ids.shape[1] :] + response = tokenizer.batch_decode( + preds, + add_special_tokens=False, + skip_special_tokens=kwargs.get("skip_special_tokens", True), + )[0] + loss = model_outputs.loss.item() + response = kwargs.get("prev_text", "") + response + + 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"{0 + 1}best_recog"] + + results = [] + response_clean = re.sub(r"[^\w\s\u3000\u4e00-\u9fff]+", "", response) + result_i = { + "key": key[0], + "text": re.sub(r"\s+", " ", response.replace("/sil", " ")), + "text_tn": response_clean, + "label": label, + } + if loss is not None: + result_i["loss"] = loss + results.append(result_i) + + for ctc_result, result in zip(ctc_results, results): + result["ctc_text"] = ctc_result["text"].replace("<|nospeech|>", "") + target_ids = torch.tensor( + self.ctc_tokenizer.encode(result["ctc_text"]), dtype=torch.int64 + ) + result["ctc_timestamps"] = forced_align( + ctc_result["ctc_logits"], target_ids, self.blank_id + ) + target_ids = torch.tensor(self.ctc_tokenizer.encode(result["text"]), dtype=torch.int64) + result["timestamps"] = forced_align(ctc_result["ctc_logits"], target_ids, self.blank_id) + for timestamps in [result["timestamps"], result["ctc_timestamps"]]: + for timestamp in timestamps: + timestamp["token"] = self.ctc_tokenizer.decode([timestamp["token"]]) + timestamp["start_time"] = timestamp["start_time"] * 6 * 10 / 1000 + timestamp["end_time"] = timestamp["end_time"] * 6 * 10 / 1000 + + if ibest_writer is not None: + ibest_writer["text"][key[0]] = response.replace("\n", " ") + ibest_writer["label"][key[0]] = label.replace("\n", " ") + ibest_writer["text_tn"][key[0]] = response_clean + + return results, meta_data + + @staticmethod + def from_pretrained(model: str = None, **kwargs): + from funasr import AutoModel + + model, kwargs = AutoModel.build_model(model=model, trust_remote_code=True, **kwargs) + + return model, kwargs diff --git a/replaced_files/klx_r200_8f/cif_predictor.py b/replaced_files/klx_r200_8f/cif_predictor.py new file mode 100644 index 0000000..31bb19b --- /dev/null +++ b/replaced_files/klx_r200_8f/cif_predictor.py @@ -0,0 +1,546 @@ +#!/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): + 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): + 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): + 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()) + 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): + 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, + ): + 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 torch.cuda.get_device_name() == 'R200-8F' and self.upsample_type != "cnn": + # kunlunxin doesn't support some ops in other two branches + self.upsample_type = "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, + ): + 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): + 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): + 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 + ): + 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): + 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, + ): + 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): + 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): + 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): + 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): + 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 diff --git a/requirements.txt b/requirements.txt index 0d2f9d2..be55be5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,4 +8,11 @@ Levenshtein ruamel.yaml nltk==3.7 pynini==2.1.6 -soundfile \ No newline at end of file +soundfile +transformers>=4.51.3 +funasr>=1.3.0 +zhconv +whisper_normalizer +pyopenjtalk-plus +compute-wer +openai-whisper \ No newline at end of file diff --git a/download_nltk_model.py b/test_scripts/download_nltk_model.py similarity index 100% rename from download_nltk_model.py rename to test_scripts/download_nltk_model.py diff --git a/start_funasr.sh b/test_scripts/start_funasr.sh similarity index 100% rename from start_funasr.sh rename to test_scripts/start_funasr.sh diff --git a/test_funasr.py b/test_scripts/test_funasr.py similarity index 100% rename from test_funasr.py rename to test_scripts/test_funasr.py