diff --git a/.gitattributes b/.gitattributes index ff0fe04b9997587d99078c0a10d96c42d0ca996b..207058e8ec1b4b07738d1398d3a20612cd827b19 100644 --- a/.gitattributes +++ b/.gitattributes @@ -46,3 +46,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text *.bmodel filter=lfs diff=lfs merge=lfs -textbmodel/tpu_perf-1.2.35-py3-none-manylinux2014_aarch64.whl filter=lfs diff=lfs merge=lfs -text +system_ref.wav filter=lfs diff=lfs merge=lfs -text diff --git a/funasr/models/paraformer_streaming/export_meta.py b/funasr/models/paraformer_streaming/export_meta.py new file mode 100644 index 0000000000000000000000000000000000000000..6815688494fcf65dafa4cf23ea4b0cb363efff9a --- /dev/null +++ b/funasr/models/paraformer_streaming/export_meta.py @@ -0,0 +1,136 @@ +#!/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 types +import torch +from funasr.register import tables + + +def export_rebuild_model(model, **kwargs): + # self.device = kwargs.get("device") + is_onnx = kwargs.get("type", "onnx") == "onnx" + encoder_class = tables.encoder_classes.get(kwargs["encoder"] + "Export") + model.encoder = encoder_class(model.encoder, onnx=is_onnx) + + predictor_class = tables.predictor_classes.get(kwargs["predictor"] + "Export") + model.predictor = predictor_class(model.predictor, onnx=is_onnx) + + if kwargs["decoder"] == "ParaformerSANMDecoder": + kwargs["decoder"] = "ParaformerSANMDecoderOnline" + decoder_class = tables.decoder_classes.get(kwargs["decoder"] + "Export") + model.decoder = decoder_class(model.decoder, onnx=is_onnx) + + from funasr.utils.torch_function import sequence_mask + + model.make_pad_mask = sequence_mask(max_seq_len=None, flip=False) + + import copy + import types + + encoder_model = copy.copy(model) + decoder_model = copy.copy(model) + + # encoder + encoder_model.forward = types.MethodType(export_encoder_forward, encoder_model) + encoder_model.export_dummy_inputs = types.MethodType(export_encoder_dummy_inputs, encoder_model) + encoder_model.export_input_names = types.MethodType(export_encoder_input_names, encoder_model) + encoder_model.export_output_names = types.MethodType(export_encoder_output_names, encoder_model) + encoder_model.export_dynamic_axes = types.MethodType(export_encoder_dynamic_axes, encoder_model) + encoder_model.export_name = "model" # types.MethodType(export_encoder_name, encoder_model) + + # decoder + decoder_model.forward = types.MethodType(export_decoder_forward, decoder_model) + decoder_model.export_dummy_inputs = types.MethodType(export_decoder_dummy_inputs, decoder_model) + decoder_model.export_input_names = types.MethodType(export_decoder_input_names, decoder_model) + decoder_model.export_output_names = types.MethodType(export_decoder_output_names, decoder_model) + decoder_model.export_dynamic_axes = types.MethodType(export_decoder_dynamic_axes, decoder_model) + decoder_model.export_name = "decoder" # types.MethodType(export_decoder_name, decoder_model) + + return encoder_model, decoder_model + + +def export_encoder_forward( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, +): + # a. To device + batch = {"speech": speech, "speech_lengths": speech_lengths, "online": True} + # batch = to_device(batch, device=self.device) + + enc, enc_len = self.encoder(**batch) + mask = self.make_pad_mask(enc_len)[:, None, :] + alphas, _ = self.predictor.forward_cnn(enc, mask) + + return enc, enc_len, alphas + + +def export_encoder_dummy_inputs(self): + speech = torch.randn(2, 30, 560) + speech_lengths = torch.tensor([6, 30], dtype=torch.int32) + return (speech, speech_lengths) + + +def export_encoder_input_names(self): + return ["speech", "speech_lengths"] + + +def export_encoder_output_names(self): + return ["enc", "enc_len", "alphas"] + + +def export_encoder_dynamic_axes(self): + return { + "speech": {0: "batch_size", 1: "feats_length"}, + "speech_lengths": { + 0: "batch_size", + }, + "enc": {0: "batch_size", 1: "feats_length"}, + "enc_len": { + 0: "batch_size", + }, + "alphas": {0: "batch_size", 1: "feats_length"}, + } + + +def export_encoder_name(self): + return "model.onnx" + + +def export_decoder_forward( + self, + enc: torch.Tensor, + enc_len: torch.Tensor, + acoustic_embeds: torch.Tensor, + acoustic_embeds_len: torch.Tensor, + *args, +): + decoder_out, out_caches = self.decoder( + enc, enc_len, acoustic_embeds, acoustic_embeds_len, *args + ) + sample_ids = decoder_out.argmax(dim=-1) + + return decoder_out, sample_ids, out_caches + + +def export_decoder_dummy_inputs(self): + dummy_inputs = self.decoder.get_dummy_inputs(enc_size=self.encoder._output_size) + return dummy_inputs + + +def export_decoder_input_names(self): + return self.decoder.get_input_names() + + +def export_decoder_output_names(self): + return self.decoder.get_output_names() + + +def export_decoder_dynamic_axes(self): + return self.decoder.get_dynamic_axes() + + +def export_decoder_name(self): + return "decoder.onnx" diff --git a/funasr/models/paraformer_streaming/model.py b/funasr/models/paraformer_streaming/model.py new file mode 100644 index 0000000000000000000000000000000000000000..a93333d8602f1d149e74bfa3a0d7076dd2c9f96d --- /dev/null +++ b/funasr/models/paraformer_streaming/model.py @@ -0,0 +1,702 @@ +#!/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 typing import Dict, Tuple +from contextlib import contextmanager +from distutils.version import LooseVersion + +from funasr.register import tables +from funasr.models.ctc.ctc import CTC +from funasr.utils import postprocess_utils +from funasr.metrics.compute_acc import th_accuracy +from funasr.utils.datadir_writer import DatadirWriter +from funasr.models.paraformer.model import Paraformer +from funasr.models.paraformer.search import Hypothesis +from funasr.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.utils.run_bmodel import EngineOV +import numpy as np + +if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): + from torch.amp import autocast +else: + # Nothing to do if torch<1.6.0 + @contextmanager + def autocast(enabled=True): + yield + + +@tables.register("model_classes", "ParaformerStreaming") +class ParaformerStreaming(Paraformer): + """ + Author: Speech Lab of DAMO Academy, Alibaba Group + Paraformer: Fast and Accurate Parallel Transformer for Non-autoregressive End-to-End Speech Recognition + https://arxiv.org/abs/2206.08317 + """ + + def __init__( + self, + *args, + **kwargs, + ): + + super().__init__(*args, **kwargs) + + self.sampling_ratio = kwargs.get("sampling_ratio", 0.2) + + self.scama_mask = None + if ( + hasattr(self.encoder, "overlap_chunk_cls") + and 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 = kwargs.get("decoder_attention_chunk_type", "chunk") + + self.encoder_bmodel = EngineOV(kwargs["model_path"]+"/encoder_fp32.bmodel", device_id=kwargs['dev_id']) + self.decoder0_bmodel = EngineOV(kwargs["model_path"]+"/decoder0_fp32.bmodel", device_id=kwargs['dev_id']) + self.decoder1_bmodel = EngineOV(kwargs["model_path"]+"/decoder1_fp32.bmodel", device_id=kwargs['dev_id']) + + self.predictor.cif_conv1d.weight = torch.load(kwargs["model_path"]+"/cif_conv1d_weight.pt",weights_only=True) + self.predictor.cif_conv1d.bias = torch.load(kwargs["model_path"]+"/cif_conv1d_bias.pt",weights_only=True) + self.predictor.cif_output.weight = torch.load(kwargs["model_path"]+"/cif_output_weight.pt",weights_only=True) + self.predictor.cif_output.bias = torch.load(kwargs["model_path"]+"/cif_output_bias.pt",weights_only=True) + + def forward( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + text: torch.Tensor, + text_lengths: torch.Tensor, + **kwargs, + ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]: + """Encoder + Decoder + Calc loss + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + text: (Batch, Length) + text_lengths: (Batch,) + """ + decoding_ind = kwargs.get("decoding_ind") + if len(text_lengths.size()) > 1: + text_lengths = text_lengths[:, 0] + if len(speech_lengths.size()) > 1: + speech_lengths = speech_lengths[:, 0] + + batch_size = speech.shape[0] + + # Encoder + if hasattr(self.encoder, "overlap_chunk_cls"): + ind = self.encoder.overlap_chunk_cls.random_choice(self.training, decoding_ind) + encoder_out, encoder_out_lens = self.encode(speech, speech_lengths, ind=ind) + else: + encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) + + loss_ctc, cer_ctc = None, None + loss_pre = None + stats = dict() + + # decoder: CTC branch + + if self.ctc_weight > 0.0: + if hasattr(self.encoder, "overlap_chunk_cls"): + encoder_out_ctc, encoder_out_lens_ctc = self.encoder.overlap_chunk_cls.remove_chunk( + encoder_out, encoder_out_lens, chunk_outs=None + ) + else: + encoder_out_ctc, encoder_out_lens_ctc = encoder_out, encoder_out_lens + + loss_ctc, cer_ctc = self._calc_ctc_loss( + encoder_out_ctc, encoder_out_lens_ctc, text, text_lengths + ) + # Collect CTC branch stats + stats["loss_ctc"] = loss_ctc.detach() if loss_ctc is not None else None + stats["cer_ctc"] = cer_ctc + + # decoder: Attention decoder branch + loss_att, acc_att, cer_att, wer_att, loss_pre, pre_loss_att = self._calc_att_predictor_loss( + encoder_out, encoder_out_lens, text, text_lengths + ) + + # 3. CTC-Att loss definition + if self.ctc_weight == 0.0: + loss = loss_att + loss_pre * self.predictor_weight + else: + loss = ( + self.ctc_weight * loss_ctc + + (1 - self.ctc_weight) * 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["pre_loss_att"] = pre_loss_att.detach() if pre_loss_att is not None else None + stats["acc"] = acc_att + stats["cer"] = cer_att + stats["wer"] = wer_att + stats["loss_pre"] = loss_pre.detach().cpu() if loss_pre is not None else None + + stats["loss"] = torch.clone(loss.detach()) + + # force_gatherable: to-device and to-tensor if scalar for DataParallel + if self.length_normalized_loss: + batch_size = (text_lengths + self.predictor_bias).sum() + loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) + return loss, stats, weight + + def encode_chunk( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + cache: dict = None, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Frontend + Encoder. Note that this method is used by asr_inference.py + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + ind: int + """ + with autocast("cuda",enabled=False): + + # Data augmentation + if self.specaug is not None and self.training: + speech, speech_lengths = self.specaug(speech, speech_lengths) + + # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN + if self.normalize is not None: + speech, speech_lengths = self.normalize(speech, speech_lengths) + + # Forward encoder + encoder_out, encoder_out_lens, _ = self.encoder.forward_chunk( + speech, speech_lengths, cache=cache["encoder"] + ) + if isinstance(encoder_out, tuple): + encoder_out = encoder_out[0] + + return encoder_out, torch.tensor([encoder_out.size(1)]) + + def _calc_att_predictor_loss( + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, + ): + encoder_out_mask = ( + ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :] + ).to(encoder_out.device) + if self.predictor_bias == 1: + _, ys_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id) + ys_pad_lens = ys_pad_lens + self.predictor_bias + 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_pad, + encoder_out_mask, + ignore_id=self.ignore_id, + mask_chunk_predictor=mask_chunk_predictor, + target_label_length=ys_pad_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_pad_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 + ) + # 0. sampler + decoder_out_1st = None + pre_loss_att = None + if self.sampling_ratio > 0.0: + + if self.use_1st_decoder_loss: + sematic_embeds, decoder_out_1st, pre_loss_att = self.sampler_with_grad( + encoder_out, + encoder_out_lens, + ys_pad, + ys_pad_lens, + pre_acoustic_embeds, + scama_mask, + ) + else: + sematic_embeds, decoder_out_1st = self.sampler( + encoder_out, + encoder_out_lens, + ys_pad, + ys_pad_lens, + pre_acoustic_embeds, + scama_mask, + ) + else: + sematic_embeds = pre_acoustic_embeds + + # 1. Forward decoder + decoder_outs = self.decoder( + encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens, scama_mask + ) + decoder_out, _ = decoder_outs[0], decoder_outs[1] + + if decoder_out_1st is None: + decoder_out_1st = decoder_out + # 2. Compute attention loss + loss_att = self.criterion_att(decoder_out, ys_pad) + acc_att = th_accuracy( + decoder_out_1st.view(-1, self.vocab_size), + ys_pad, + ignore_label=self.ignore_id, + ) + loss_pre = self.criterion_pre(ys_pad_lens.type_as(pre_token_length), pre_token_length) + + # Compute cer/wer using attention-decoder + if self.training or self.error_calculator is None: + cer_att, wer_att = None, None + else: + ys_hat = decoder_out_1st.argmax(dim=-1) + cer_att, wer_att = self.error_calculator(ys_hat.cpu(), ys_pad.cpu()) + + return loss_att, acc_att, cer_att, wer_att, loss_pre, pre_loss_att + + def sampler( + self, + encoder_out, + encoder_out_lens, + ys_pad, + ys_pad_lens, + pre_acoustic_embeds, + chunk_mask=None, + ): + + tgt_mask = (~make_pad_mask(ys_pad_lens, maxlen=ys_pad_lens.max())[:, :, None]).to( + ys_pad.device + ) + ys_pad_masked = ys_pad * tgt_mask[:, :, 0] + if self.share_embedding: + ys_pad_embed = self.decoder.output_layer.weight[ys_pad_masked] + else: + ys_pad_embed = self.decoder.embed(ys_pad_masked) + with torch.no_grad(): + decoder_outs = self.decoder( + encoder_out, encoder_out_lens, pre_acoustic_embeds, ys_pad_lens, chunk_mask + ) + decoder_out, _ = decoder_outs[0], decoder_outs[1] + pred_tokens = decoder_out.argmax(-1) + nonpad_positions = ys_pad.ne(self.ignore_id) + seq_lens = (nonpad_positions).sum(1) + same_num = ((pred_tokens == ys_pad) & nonpad_positions).sum(1) + input_mask = torch.ones_like(nonpad_positions) + bsz, seq_len = ys_pad.size() + for li in range(bsz): + target_num = ( + ((seq_lens[li] - same_num[li].sum()).float()) * self.sampling_ratio + ).long() + if target_num > 0: + input_mask[li].scatter_( + dim=0, index=torch.randperm(seq_lens[li])[:target_num].cuda(), value=0 + ) + input_mask = input_mask.eq(1) + input_mask = input_mask.masked_fill(~nonpad_positions, False) + input_mask_expand_dim = input_mask.unsqueeze(2).to(pre_acoustic_embeds.device) + + sematic_embeds = pre_acoustic_embeds.masked_fill( + ~input_mask_expand_dim, 0 + ) + ys_pad_embed.masked_fill(input_mask_expand_dim, 0) + return sematic_embeds * tgt_mask, decoder_out * tgt_mask + + def calc_predictor(self, encoder_out, encoder_out_lens): + + encoder_out_mask = ( + ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :] + ).to(encoder_out.device) + 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, pre_peak_index = self.predictor( + encoder_out, + None, + encoder_out_mask, + ignore_id=self.ignore_id, + mask_chunk_predictor=mask_chunk_predictor, + target_label_length=None, + ) + predictor_alignments, predictor_alignments_len = self.predictor.gen_frame_alignments( + pre_alphas, + encoder_out_lens + 1 if self.predictor.tail_threshold > 0.0 else 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=None, + is_training=self.training, + ) + self.scama_mask = scama_mask + + return pre_acoustic_embeds, pre_token_length, pre_alphas, pre_peak_index + + def calc_predictor_chunk(self, encoder_out, encoder_out_lens, cache=None, **kwargs): + is_final = kwargs.get("is_final", False) + + return self.predictor.forward_chunk(encoder_out, cache["encoder"], is_final=is_final) + + def cal_decoder_with_predictor( + self, encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens + ): + decoder_outs = self.decoder( + encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens, self.scama_mask + ) + decoder_out = decoder_outs[0] + decoder_out = torch.log_softmax(decoder_out, dim=-1) + return decoder_out, ys_pad_lens + + def cal_decoder_with_predictor_chunk( + self, encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens, cache=None + ): + decoder_outs = self.decoder.forward_chunk(encoder_out, sematic_embeds, cache["decoder"]) + decoder_out = decoder_outs + decoder_out = torch.log_softmax(decoder_out, dim=-1) + return decoder_out, ys_pad_lens + + def init_cache(self, cache: dict = {}, **kwargs): + chunk_size = kwargs.get("chunk_size", [0, 10, 5]) + encoder_chunk_look_back = kwargs.get("encoder_chunk_look_back", 0) + decoder_chunk_look_back = kwargs.get("decoder_chunk_look_back", 0) + batch_size = 1 + + enc_output_size = kwargs["encoder_conf"]["output_size"] + feats_dims = kwargs["frontend_conf"]["n_mels"] * kwargs["frontend_conf"]["lfr_m"] + cache_encoder = { + "start_idx": 0, + "cif_hidden": torch.zeros((batch_size, 1, enc_output_size)), + "cif_alphas": torch.zeros((batch_size, 1)), + "chunk_size": chunk_size, + "encoder_chunk_look_back": encoder_chunk_look_back, + "last_chunk": False, + "opt": None, + "feats": torch.zeros((batch_size, chunk_size[0] + chunk_size[2], feats_dims)), + "tail_chunk": False, + } + cache["encoder"] = cache_encoder + + cache_decoder = { + "decode_fsmn": None, + "decoder_chunk_look_back": decoder_chunk_look_back, + "opt": None, + "chunk_size": chunk_size, + } + cache["decoder"] = cache_decoder + cache["frontend"] = {} + cache["prev_samples"] = torch.empty(0) + + return cache + + def generate_chunk( + self, + speech, + speech_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): + cache = kwargs.get("cache", {}) + speech = speech.to(device=kwargs["device"]) + speech_lengths = speech_lengths.to(device=kwargs["device"]) + + # Encoder + """ + encoder_out, encoder_out_lens = self.encode_chunk( + speech, speech_lengths, cache=cache, is_final=kwargs.get("is_final", False) + ) + if isinstance(encoder_out, tuple): + encoder_out = encoder_out[0] + """ + speech = speech.detach().numpy() + start_idx = (torch.tensor([0], dtype=torch.int64).detach().numpy()).astype(np.int32) + chunk_size = (torch.tensor(cache['encoder']['chunk_size'], dtype=torch.int64).detach().numpy()).astype(np.int32) + feats = cache['encoder']['feats'].detach().numpy() + outputs = self.encoder_bmodel([speech, start_idx, chunk_size, feats]) + cache['encoder']['start_idx'] = torch.from_numpy(outputs[1]) + cache['encoder']['feats'] = torch.from_numpy(outputs[2]) + encoder_out = torch.from_numpy(outputs[0]) + encoder_out_lens = torch.tensor([encoder_out.size(1)]) + + # predictor + predictor_outs = self.calc_predictor_chunk( + encoder_out, encoder_out_lens, cache=cache, is_final=kwargs.get("is_final", False) + ) + pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = ( + predictor_outs[0], + predictor_outs[1], + predictor_outs[2], + predictor_outs[3], + ) + pre_token_length = pre_token_length.round().long() + if torch.max(pre_token_length) < 1: + return [] + """ + decoder_outs = self.cal_decoder_with_predictor_chunk( + encoder_out, encoder_out_lens, pre_acoustic_embeds, pre_token_length, cache=cache + ) + decoder_out, ys_pad_lens = decoder_outs[0], decoder_outs[1] + """ + encoder_out = encoder_out.detach().numpy() + pre_acoustic_embeds = pre_acoustic_embeds.detach().numpy() + lens = (pre_token_length.detach().numpy()).astype(np.int32) + input = [encoder_out, pre_acoustic_embeds] + if cache['decoder']['decode_fsmn'] is None: + outputs = self.decoder0_bmodel(input) + else: + input.append(lens) + for i in range(16): + input.append(cache['decoder']['decode_fsmn'][i].detach().numpy()) + outputs = self.decoder1_bmodel(input) + decode_fsmn = [] + for i in range(16): + decode_fsmn.append(torch.from_numpy(outputs[i+1])) + cache['decoder']['decode_fsmn'] = decode_fsmn + decoder_out = torch.from_numpy(outputs[0]) + + results = [] + b, n, d = decoder_out.size() + if isinstance(key[0], (list, tuple)): + key = key[0] + for i in range(b): + x = encoder_out[i, : encoder_out_lens[i], :] + am_scores = decoder_out[i, : pre_token_length[i], :] + if self.beam_search is not None: + nbest_hyps = self.beam_search( + x=x, + am_scores=am_scores, + maxlenratio=kwargs.get("maxlenratio", 0.0), + minlenratio=kwargs.get("minlenratio", 0.0), + ) + + nbest_hyps = nbest_hyps[: self.nbest] + else: + + yseq = am_scores.argmax(dim=-1) + score = am_scores.max(dim=-1)[0] + score = torch.sum(score, dim=-1) + # pad with mask tokens to ensure compatibility with sos/eos tokens + yseq = torch.tensor([self.sos] + yseq.tolist() + [self.eos], device=yseq.device) + nbest_hyps = [Hypothesis(yseq=yseq, score=score)] + for nbest_idx, hyp in enumerate(nbest_hyps): + + # remove sos/eos and get results + last_pos = -1 + if isinstance(hyp.yseq, list): + token_int = hyp.yseq[1:last_pos] + else: + token_int = hyp.yseq[1:last_pos].tolist() + + # remove blank symbol id, which is assumed to be 0 + token_int = list( + filter( + lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int + ) + ) + + # Change integer-ids to tokens + token = tokenizer.ids2tokens(token_int) + # text = tokenizer.tokens2text(token) + + result_i = token + + results.extend(result_i) + + return results + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + cache: dict = {}, + **kwargs, + ): + + # init beamsearch + is_use_ctc = kwargs.get("decoding_ctc_weight", 0.0) > 0.00001 and self.ctc != None + is_use_lm = ( + kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None + ) + if self.beam_search is None and (is_use_lm or is_use_ctc): + logging.info("enable beam_search") + self.init_beam_search(**kwargs) + self.nbest = kwargs.get("nbest", 1) + + if len(cache) == 0: + self.init_cache(cache, **kwargs) + + meta_data = {} + chunk_size = kwargs.get("chunk_size", [0, 10, 5]) + chunk_stride_samples = int(chunk_size[1] * 960) # 600ms + + time1 = time.perf_counter() + cfg = {"is_final": kwargs.get("is_final", False)} + 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, + cache=cfg, + ) + _is_final = cfg["is_final"] # if data_in is a file or url, set is_final=True + + time2 = time.perf_counter() + meta_data["load_data"] = f"{time2 - time1:0.3f}" + assert len(audio_sample_list) == 1, "batch_size must be set 1" + + audio_sample = torch.cat((cache["prev_samples"], audio_sample_list[0])) + + n = int(len(audio_sample) // chunk_stride_samples + int(_is_final)) + m = int(len(audio_sample) % chunk_stride_samples * (1 - int(_is_final))) + tokens = [] + for i in range(n): + kwargs["is_final"] = _is_final and i == n - 1 + audio_sample_i = audio_sample[i * chunk_stride_samples : (i + 1) * chunk_stride_samples] + if kwargs["is_final"] and len(audio_sample_i) < 960: + cache["encoder"]["tail_chunk"] = True + speech = cache["encoder"]["feats"] + speech_lengths = torch.tensor([speech.shape[1]], dtype=torch.int64).to( + speech.device + ) + else: + # extract fbank feats + speech, speech_lengths = extract_fbank( + [audio_sample_i], + data_type=kwargs.get("data_type", "sound"), + frontend=frontend, + cache=cache["frontend"], + is_final=kwargs["is_final"], + ) + 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 + ) + + tokens_i = self.generate_chunk( + speech, + speech_lengths, + key=key, + tokenizer=tokenizer, + cache=cache, + frontend=frontend, + **kwargs, + ) + tokens.extend(tokens_i) + + text_postprocessed, _ = postprocess_utils.sentence_postprocess(tokens) + + result_i = {"key": key[0], "text": text_postprocessed} + result = [result_i] + + cache["prev_samples"] = audio_sample[:-m] + if _is_final: + self.init_cache(cache, **kwargs) + + if kwargs.get("output_dir"): + if not hasattr(self, "writer"): + self.writer = DatadirWriter(kwargs.get("output_dir")) + ibest_writer = self.writer[f"{1}best_recog"] + ibest_writer["token"][key[0]] = " ".join(tokens) + ibest_writer["text"][key[0]] = text_postprocessed + + return result, meta_data + + def export(self, **kwargs): + from .export_meta import export_rebuild_model + + models = export_rebuild_model(model=self, **kwargs) + return models diff --git a/funasr/models/paraformer_streaming/template.yaml b/funasr/models/paraformer_streaming/template.yaml new file mode 100644 index 0000000000000000000000000000000000000000..889971ad12e93edcbbc3dccd6f20c03051448588 --- /dev/null +++ b/funasr/models/paraformer_streaming/template.yaml @@ -0,0 +1,143 @@ +# This is an example that demonstrates how to configure a model file. +# You can modify the configuration according to your own requirements. + +# to print the register_table: +# from funasr.register import tables +# tables.print() + +# network architecture +model: ParaformerStreaming +model_conf: + ctc_weight: 0.0 + lsm_weight: 0.1 + length_normalized_loss: true + predictor_weight: 1.0 + predictor_bias: 1 + sampling_ratio: 0.75 + +# encoder +encoder: SANMEncoderChunkOpt +encoder_conf: + output_size: 512 + attention_heads: 4 + linear_units: 2048 + num_blocks: 50 + dropout_rate: 0.1 + positional_dropout_rate: 0.1 + attention_dropout_rate: 0.1 + input_layer: pe_online + pos_enc_class: SinusoidalPositionEncoder + normalize_before: true + kernel_size: 11 + sanm_shfit: 0 + selfattention_layer_type: sanm + chunk_size: + - 12 + - 15 + stride: + - 8 + - 10 + pad_left: + - 0 + - 0 + encoder_att_look_back_factor: + - 4 + - 4 + decoder_att_look_back_factor: + - 1 + - 1 + +# decoder +decoder: ParaformerSANMDecoder +decoder_conf: + attention_heads: 4 + linear_units: 2048 + num_blocks: 16 + dropout_rate: 0.1 + positional_dropout_rate: 0.1 + self_attention_dropout_rate: 0.1 + src_attention_dropout_rate: 0.1 + att_layer_num: 16 + kernel_size: 11 + sanm_shfit: 5 + +predictor: CifPredictorV2 +predictor_conf: + idim: 512 + threshold: 1.0 + l_order: 1 + r_order: 1 + tail_threshold: 0.45 + +# frontend related +frontend: WavFrontendOnline +frontend_conf: + fs: 16000 + window: hamming + n_mels: 80 + frame_length: 25 + frame_shift: 10 + lfr_m: 7 + lfr_n: 6 + +specaug: SpecAugLFR +specaug_conf: + apply_time_warp: false + time_warp_window: 5 + time_warp_mode: bicubic + apply_freq_mask: true + freq_mask_width_range: + - 0 + - 30 + lfr_rate: 6 + num_freq_mask: 1 + apply_time_mask: true + time_mask_width_range: + - 0 + - 12 + num_time_mask: 1 + +train_conf: + accum_grad: 1 + grad_clip: 5 + max_epoch: 150 + val_scheduler_criterion: + - valid + - acc + best_model_criterion: + - - valid + - acc + - max + keep_nbest_models: 10 + log_interval: 50 + +optim: adam +optim_conf: + lr: 0.0005 +scheduler: warmuplr +scheduler_conf: + warmup_steps: 30000 + +dataset: AudioDataset +dataset_conf: + index_ds: IndexDSJsonl + batch_sampler: BatchSampler + batch_type: example # example or length + batch_size: 1 # if batch_type is example, batch_size is the numbers of samples; if length, batch_size is source_token_len+target_token_len; + max_token_length: 2048 # filter samples if source_token_len+target_token_len > max_token_length, + buffer_size: 500 + shuffle: True + num_workers: 0 + +tokenizer: CharTokenizer +tokenizer_conf: + unk_symbol: + split_with_space: true + + +ctc_conf: + dropout_rate: 0.0 + ctc_type: builtin + reduce: true + ignore_nan_grad: true +normalize: null diff --git a/funasr/models/qwen_audio/__init__.py b/funasr/models/qwen_audio/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/models/qwen_audio/audio.py b/funasr/models/qwen_audio/audio.py new file mode 100644 index 0000000000000000000000000000000000000000..a1aad79ea60955839b7b37c2891ec405ce798de1 --- /dev/null +++ b/funasr/models/qwen_audio/audio.py @@ -0,0 +1,442 @@ +import base64 +import gzip +from dataclasses import dataclass +from typing import Dict, Iterable, Optional, List + +import numpy as np +import torch +import torch.nn.functional as F +from torch import Tensor, nn +from subprocess import CalledProcessError, run, Popen, PIPE + +import os +from functools import lru_cache +from typing import Optional, Union +from funasr.register import tables + + +def exact_div(x, y): + assert x % y == 0 + return x // y + + +# hard-coded audio hyperparameters +SAMPLE_RATE = 16000 +N_FFT = 400 +N_MELS = 80 +HOP_LENGTH = 160 +CHUNK_LENGTH = 30 +N_SAMPLES = CHUNK_LENGTH * SAMPLE_RATE # 480000 samples in a 30-second chunk +N_FRAMES = exact_div(N_SAMPLES, HOP_LENGTH) # 3000 frames in a mel spectrogram input + +N_SAMPLES_PER_TOKEN = HOP_LENGTH * 2 # the initial convolutions has stride 2 +FRAMES_PER_SECOND = exact_div(SAMPLE_RATE, HOP_LENGTH) # 10ms per audio frame +TOKENS_PER_SECOND = exact_div(SAMPLE_RATE, N_SAMPLES_PER_TOKEN) # 20ms per audio token + + +def get_T_after_cnn(L_in, dilation=1): + for padding, kernel_size, stride in eval("[(1,3,1)] + [(1,3,2)] "): + L_out = L_in + 2 * padding - dilation * (kernel_size - 1) - 1 + L_out = 1 + L_out // stride + L_in = L_out + return L_out + + +def load_bytesio_audio(content, sr: int = SAMPLE_RATE): + cmd = [ + "ffmpeg", + "-nostdin", + "-threads", + "0", + "-i", + "pipe:", + "-f", + "s16le", + "-ac", + "1", + "-acodec", + "pcm_s16le", + "-ar", + str(sr), + "pipe:", + ] + p = Popen(cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE, bufsize=-1) + out, _ = p.communicate(input=content) + return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0 + + +def load_audio(file: str, sr: int = SAMPLE_RATE): + """ + Open an audio file and read as mono waveform, resampling as necessary + + Parameters + ---------- + file: str + The audio file to open + + sr: int + The sample rate to resample the audio if necessary + + Returns + ------- + A NumPy array containing the audio waveform, in float32 dtype. + """ + + # This launches a subprocess to decode audio while down-mixing + # and resampling as necessary. Requires the ffmpeg CLI in PATH. + # fmt: off + cmd = [ + "ffmpeg", + "-nostdin", + "-threads", "0", + "-i", file, + "-f", "s16le", + "-ac", "1", + "-acodec", "pcm_s16le", + "-ar", str(sr), + "-" + ] + # fmt: on + try: + out = run(cmd, capture_output=True, check=True).stdout + except CalledProcessError as e: + raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e + + return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0 + + +def pad_or_trim(array, length: int = N_SAMPLES, *, axis: int = -1): + """ + Pad or trim the audio array to N_SAMPLES, as expected by the encoder. + """ + if torch.is_tensor(array): + if array.shape[axis] > length: + array = array.index_select(dim=axis, index=torch.arange(length, device=array.device)) + + if array.shape[axis] < length: + pad_widths = [(0, 0)] * array.ndim + pad_widths[axis] = (0, length - array.shape[axis]) + array = F.pad(array, [pad for sizes in pad_widths[::-1] for pad in sizes]) + else: + if array.shape[axis] > length: + array = array.take(indices=range(length), axis=axis) + + if array.shape[axis] < length: + pad_widths = [(0, 0)] * array.ndim + pad_widths[axis] = (0, length - array.shape[axis]) + array = np.pad(array, pad_widths) + + return array + + +def trim(array, length: int = N_SAMPLES, *, axis: int = -1): + """ + Pad or trim the audio array to N_SAMPLES, as expected by the encoder. + """ + if torch.is_tensor(array): + if array.shape[axis] > length: + array = array.index_select(dim=axis, index=torch.arange(length, device=array.device)) + else: + if array.shape[axis] > length: + array = array.take(indices=range(length), axis=axis) + return array + + +@lru_cache(maxsize=None) +def mel_filters(device, n_mels: int = N_MELS) -> torch.Tensor: + """ + load the mel filterbank matrix for projecting STFT into a Mel spectrogram. + Allows decoupling librosa dependency; saved using: + + np.savez_compressed( + "mel_filters.npz", + mel_80=librosa.filters.mel(sr=16000, n_fft=400, n_mels=80), + ) + """ + assert n_mels == 80, f"Unsupported n_mels: {n_mels}" + with np.load( + os.path.join(os.path.dirname(__file__), "mel_filters.npz") # todo + # os.path.join("assets", "mel_filters.npz") + ) as f: + return torch.from_numpy(f[f"mel_{n_mels}"]).to(device) + + +def log_mel_spectrogram( + audio: Union[str, np.ndarray, torch.Tensor], + n_mels: int = N_MELS, + padding: int = 0, + device: Optional[Union[str, torch.device]] = None, +): + """ + Compute the log-Mel spectrogram of + + Parameters + ---------- + audio: Union[str, np.ndarray, torch.Tensor], shape = (*) + The path to audio or either a NumPy array or Tensor containing the audio waveform in 16 kHz + + n_mels: int + The number of Mel-frequency filters, only 80 is supported + + padding: int + Number of zero samples to pad to the right + + device: Optional[Union[str, torch.device]] + If given, the audio tensor is moved to this device before STFT + + Returns + ------- + torch.Tensor, shape = (80, n_frames) + A Tensor that contains the Mel spectrogram + """ + if not torch.is_tensor(audio): + if isinstance(audio, str): + audio = load_audio(audio) + audio = torch.from_numpy(audio) + + if device is not None: + audio = audio.to(device) + if padding > 0: + audio = F.pad(audio, (0, padding)) + window = torch.hann_window(N_FFT).to(audio.device) + stft = torch.stft(audio, N_FFT, HOP_LENGTH, window=window, return_complex=True) + magnitudes = stft[..., :-1].abs() ** 2 + + filters = mel_filters(audio.device, n_mels) + mel_spec = filters @ magnitudes + + log_spec = torch.clamp(mel_spec, min=1e-10).log10() + log_spec = torch.maximum(log_spec, log_spec.max() - 8.0) + log_spec = (log_spec + 4.0) / 4.0 + return log_spec + + +@dataclass +class ModelDimensions: + n_mels: int + n_audio_ctx: int + n_audio_state: int + n_audio_head: int + n_audio_layer: int + n_vocab: int + n_text_ctx: int + n_text_state: int + n_text_head: int + n_text_layer: int + + +class LayerNorm(nn.LayerNorm): + def forward(self, x: Tensor) -> Tensor: + # return super().forward(x.float()).type(x.dtype) + return super().forward(x).type(x.dtype) + + +class Linear(nn.Linear): + def forward(self, x: Tensor) -> Tensor: + return F.linear( + x, + self.weight.to(x.dtype), + None if self.bias is None else self.bias.to(x.dtype), + ) + + +class Conv1d(nn.Conv1d): + def _conv_forward(self, x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> Tensor: + return super()._conv_forward( + x, weight.to(x.dtype), None if bias is None else bias.to(x.dtype) + ) + + +def sinusoids(length, channels, max_timescale=10000): + """Returns sinusoids for positional embedding""" + assert channels % 2 == 0 + log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1) + inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2)) + scaled_time = torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :] + return torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1) + + +class MultiHeadAttention(nn.Module): + def __init__(self, n_state: int, n_head: int): + super().__init__() + self.n_head = n_head + self.query = Linear(n_state, n_state) + self.key = Linear(n_state, n_state, bias=False) + self.value = Linear(n_state, n_state) + self.out = Linear(n_state, n_state) + + def forward( + self, + x: Tensor, + xa: Optional[Tensor] = None, + mask: Optional[Tensor] = None, + kv_cache: Optional[dict] = None, + ): + q = self.query(x) + + if kv_cache is None or xa is None or self.key not in kv_cache: + # hooks, if installed (i.e. kv_cache is not None), will prepend the cached kv tensors; + # otherwise, perform key/value projections for self- or cross-attention as usual. + k = self.key(x if xa is None else xa) + v = self.value(x if xa is None else xa) + else: + # for cross-attention, calculate keys and values once and reuse in subsequent calls. + k = kv_cache[self.key] + v = kv_cache[self.value] + + wv, qk = self.qkv_attention(q, k, v, mask) + return self.out(wv), qk + + def qkv_attention(self, q: Tensor, k: Tensor, v: Tensor, mask: Optional[Tensor] = None): + n_batch, n_ctx, n_state = q.shape + scale = (n_state // self.n_head) ** -0.25 + q = q.view(*q.shape[:2], self.n_head, -1).permute(0, 2, 1, 3) * scale + k = k.view(*k.shape[:2], self.n_head, -1).permute(0, 2, 3, 1) * scale + v = v.view(*v.shape[:2], self.n_head, -1).permute(0, 2, 1, 3) + + qk = q @ k + if mask is not None: + qk += mask + + w = F.softmax(qk, dim=-1).to(q.dtype) + return (w @ v).permute(0, 2, 1, 3).flatten(start_dim=2), qk.detach() + + +class ResidualAttentionBlock(nn.Module): + def __init__(self, n_state: int, n_head: int, cross_attention: bool = False): + super().__init__() + + self.attn = MultiHeadAttention(n_state, n_head) + self.attn_ln = LayerNorm(n_state) + + self.cross_attn = MultiHeadAttention(n_state, n_head) if cross_attention else None + self.cross_attn_ln = LayerNorm(n_state) if cross_attention else None + + n_mlp = n_state * 4 + self.mlp = nn.Sequential(Linear(n_state, n_mlp), nn.GELU(), Linear(n_mlp, n_state)) + self.mlp_ln = LayerNorm(n_state) + + def forward( + self, + x: Tensor, + xa: Optional[Tensor] = None, + mask: Optional[Tensor] = None, + kv_cache: Optional[dict] = None, + ): + x = x + self.attn(self.attn_ln(x), mask=mask, kv_cache=kv_cache)[0] + if self.cross_attn: + x = x + self.cross_attn(self.cross_attn_ln(x), xa, kv_cache=kv_cache)[0] + x = x + self.mlp(self.mlp_ln(x)) + return x + + +@tables.register("encoder_classes", "QwenAudioEncoder") +class QwenAudioEncoder(nn.Module): + def __init__( + self, + n_mels: int, + n_ctx: int, + n_state: int, + n_head: int, + n_layer: int, + output_dim: int = 512, + avg_pool: bool = True, + add_audio_bos_eos_token: bool = True, + **kwargs, + ): + super().__init__() + self.conv1 = Conv1d(n_mels, n_state, kernel_size=3, padding=1) + self.conv2 = Conv1d(n_state, n_state, kernel_size=3, stride=2, padding=1) + self.register_buffer("positional_embedding", sinusoids(n_ctx, n_state)) + + self.blocks: Iterable[ResidualAttentionBlock] = nn.ModuleList( + [ResidualAttentionBlock(n_state, n_head) for _ in range(n_layer)] + ) + self.ln_post = LayerNorm(n_state) + + if avg_pool: + self.avg_pooler = nn.AvgPool1d(2, stride=2) + else: + self.avg_pooler = None + self.proj = nn.Linear(n_state, output_dim) + if add_audio_bos_eos_token: + self.audio_bos_eos_token = nn.Embedding(2, output_dim) + else: + self.audio_bos_eos_token = None + self.output_dim = output_dim + self.n_head = n_head + + def forward(self, x: Tensor, padding_mask: Tensor = None, audio_lengths: Tensor = None): + """ + x : torch.Tensor, shape = (batch_size, n_mels, n_ctx) + the mel spectrogram of the audio + """ + x = x.to(dtype=self.conv1.weight.dtype, device=self.conv1.weight.device) + if audio_lengths is not None: + input_mel_len = audio_lengths[:, 0] * 2 + max_mel_len_in_batch = input_mel_len.max() + x = x[:, :, :max_mel_len_in_batch] + x = F.gelu(self.conv1(x)) + x = F.gelu(self.conv2(x)) + x = x.permute(0, 2, 1) # B, L, D + bsz = x.size(0) + src_len = x.size(1) + + self.input_positional_embedding = self.positional_embedding[:src_len] + assert ( + x.shape[1:] == self.input_positional_embedding.shape + ), f"incorrect audio shape: {x.shape[1:], self.input_positional_embedding.shape}" + x = (x + self.input_positional_embedding).to(x.dtype) + if padding_mask is not None: + padding_mask = padding_mask.to( + dtype=self.conv1.weight.dtype, device=self.conv1.weight.device + ) + batch_src_len = padding_mask.size(1) + x = x[:, :batch_src_len, :] + padding_mask = padding_mask.view(bsz, -1, batch_src_len) + padding_mask_ = padding_mask.all(1) + x[padding_mask_] = 0 + key_padding_mask = ( + padding_mask_.view(bsz, 1, 1, batch_src_len) + .expand(-1, self.n_head, -1, -1) + .reshape(bsz, self.n_head, 1, batch_src_len) + ) + new_padding_mask = torch.zeros_like(key_padding_mask, dtype=x.dtype) + padding_mask = new_padding_mask.masked_fill(key_padding_mask, float("-inf")) + + for block in self.blocks: + x = block(x, mask=padding_mask) + + if self.avg_pooler: + x = x.permute(0, 2, 1) + x = self.avg_pooler(x) + x = x.permute(0, 2, 1) + + x = self.ln_post(x) + x = self.proj(x) + + if self.audio_bos_eos_token is not None: + bos = self.audio_bos_eos_token.weight[0][None, :] + eos = self.audio_bos_eos_token.weight[1][None, :] + else: + bos, eos = None, None + return x, bos, eos + + def encode(self, input_audios: Tensor, input_audio_lengths: Tensor, audio_span_tokens: List): + real_input_audio_lens = input_audio_lengths[:, 0].tolist() + max_len_in_batch = max(real_input_audio_lens) + padding_mask = torch.ones([input_audios.size(0), max_len_in_batch]).to( + dtype=self.conv1.weight.dtype, device=self.conv1.weight.device + ) + for index in range(len(input_audios)): + padding_mask[index, : input_audio_lengths[index][0].item()] = 0 + x, bos, eos = self(input_audios, padding_mask, input_audio_lengths) + output_audios = [] + for i in range(len(audio_span_tokens)): + audio_span = audio_span_tokens[i] + audio = x[i][: audio_span - 2] + if bos is not None: + audio = torch.concat([bos, audio, eos]) + assert len(audio) == audio_span + output_audios.append(audio) + return output_audios diff --git a/funasr/models/qwen_audio/model.py b/funasr/models/qwen_audio/model.py new file mode 100644 index 0000000000000000000000000000000000000000..b0af4563c6e75e7d9b1ae8ca493ffe06d3145ecd --- /dev/null +++ b/funasr/models/qwen_audio/model.py @@ -0,0 +1,151 @@ +from dataclasses import dataclass +from typing import Dict +from typing import Iterable, Optional +import time +import numpy as np +import torch +import torch.nn.functional as F +from torch import Tensor +from torch import nn +import whisper +from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank + + +from funasr.register import tables + + +@tables.register("model_classes", "Qwen/Qwen-Audio") +@tables.register("model_classes", "Qwen-Audio") +@tables.register("model_classes", "Qwen/QwenAudio") +@tables.register("model_classes", "QwenAudio") +@tables.register("model_classes", "QwenAudioWarp") +class QwenAudioWarp(nn.Module): + """ + Qwen-Audio: Advancing Universal Audio Understanding via Unified Large-Scale Audio-Language Models + https://arxiv.org/abs/2311.07919 + Modified from https://github.com/QwenLM/Qwen-Audio + """ + + def __init__(self, *args, **kwargs): + super().__init__() + from transformers import AutoModelForCausalLM, AutoTokenizer + from transformers.generation import GenerationConfig + + model_or_path = kwargs.get("model_path", "QwenAudio") + model = AutoModelForCausalLM.from_pretrained( + model_or_path, device_map="cpu", trust_remote_code=True + ) + tokenizer = AutoTokenizer.from_pretrained(model_or_path, trust_remote_code=True) + + self.model = model + self.tokenizer = tokenizer + + def forward( + self, + ): + pass + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): + if kwargs.get("batch_size", 1) > 1: + raise NotImplementedError("batch decoding is not implemented") + + meta_data = {} + # meta_data["batch_data_time"] = -1 + prompt = kwargs.get( + "prompt", "<|startoftranscription|><|en|><|transcribe|><|en|><|notimestamps|><|wo_itn|>" + ) + query = f"{prompt}" + audio_info = self.tokenizer.process_audio(query) + inputs = self.tokenizer(query, return_tensors="pt", audio_info=audio_info) + inputs = inputs.to(self.model.device) + pred = self.model.generate(**inputs, audio_info=audio_info) + response = self.tokenizer.decode( + pred.cpu()[0], skip_special_tokens=False, audio_info=audio_info + ) + + results = [] + result_i = {"key": key[0], "text": response} + + results.append(result_i) + + return results, meta_data + + +@tables.register("model_classes", "Qwen/Qwen-Audio-Chat") +@tables.register("model_classes", "Qwen/QwenAudioChat") +@tables.register("model_classes", "Qwen-Audio-Chat") +@tables.register("model_classes", "QwenAudioChat") +@tables.register("model_classes", "QwenAudioChatWarp") +class QwenAudioChatWarp(nn.Module): + def __init__(self, *args, **kwargs): + """ + Qwen-Audio: Advancing Universal Audio Understanding via Unified Large-Scale Audio-Language Models + https://arxiv.org/abs/2311.07919 + Modified from https://github.com/QwenLM/Qwen-Audio + """ + super().__init__() + from transformers import AutoModelForCausalLM, AutoTokenizer + from transformers.generation import GenerationConfig + + model_or_path = kwargs.get("model_path", "QwenAudio") + bf16 = kwargs.get("bf16", False) + fp16 = kwargs.get("fp16", False) + model = AutoModelForCausalLM.from_pretrained( + model_or_path, device_map="cpu", bf16=bf16, fp16=fp16, trust_remote_code=True + ) + tokenizer = AutoTokenizer.from_pretrained(model_or_path, trust_remote_code=True) + + self.model = model + self.tokenizer = tokenizer + + def forward( + self, + ): + pass + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): + if kwargs.get("batch_size", 1) > 1: + raise NotImplementedError("batch decoding is not implemented") + + meta_data = {} + + prompt = kwargs.get("prompt", "what does the person say?") + cache = kwargs.get("cache", {}) + history = cache.get("history", None) + if data_in[0] is not None: + # 1st dialogue turn + query = self.tokenizer.from_list_format( + [ + {"audio": data_in[0]}, # Either a local path or an url + {"text": prompt}, + ] + ) + else: + query = prompt + response, history = self.model.chat(self.tokenizer, query=query, history=history) + cache["history"] = history + # print(response) + # The person says: "mister quilter is the apostle of the middle classes and we are glad to welcome his gospel". + + results = [] + result_i = {"key": key[0], "text": response} + + results.append(result_i) + + return results, meta_data diff --git a/funasr/models/qwen_audio/template.yaml b/funasr/models/qwen_audio/template.yaml new file mode 100644 index 0000000000000000000000000000000000000000..40b902c716c531521bfc6715afcea31a34a6a85d --- /dev/null +++ b/funasr/models/qwen_audio/template.yaml @@ -0,0 +1,46 @@ +# This is an example that demonstrates how to configure a model file. +# You can modify the configuration according to your own requirements. + +# to print the register_table: +# from funasr.register import tables +# tables.print() + +# network architecture +model: WhisperWarp +model_conf: + lsm_weight: 0.1 + length_normalized_loss: true + hub: funasr # openai + init_param_path: null # large-v2 or large-v3 if hub == "openai" + + + +# only use for hub == funasr, +# if hub == openai, whisper_dims is automaticall download +whisper_dims: + 'n_mels': 80 + 'n_vocab': 51865 + 'n_audio_ctx': 1500 + 'n_audio_state': 1280 + 'n_audio_head': 20 + 'n_audio_layer': 32 + 'n_text_ctx': 448 + 'n_text_state': 1280 + 'n_text_head': 20 + 'n_text_layer': 32 + +# frontend related +frontend: WhisperFrontend +frontend_conf: + fs: 16000 + n_mels: 80 + do_pad_trim: true + +tokenizer: WhisperTokenizer +tokenizer_conf: + language: null + task: transcribe + is_multilingual: true + num_languages: 99 + +scope_map: ['none', "model."] diff --git a/funasr/models/rwkv_bat/__init__.py b/funasr/models/rwkv_bat/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/models/rwkv_bat/cuda_decoder/wkv_cuda.cu b/funasr/models/rwkv_bat/cuda_decoder/wkv_cuda.cu new file mode 100644 index 0000000000000000000000000000000000000000..1dcbc7141dc69655534214f6df642731c8c1260f --- /dev/null +++ b/funasr/models/rwkv_bat/cuda_decoder/wkv_cuda.cu @@ -0,0 +1,135 @@ +// Copied from https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v4/cuda/wkv_cuda.cu + +#include +#include + +#define MIN_VALUE (-1e38) + +template +__global__ void kernel_forward(const int B, const int T, const int C, + const F *__restrict__ const _w, const F *__restrict__ const _u, const F *__restrict__ const _k, const F *__restrict__ const _v, + F *__restrict__ const _y) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + const int _b = idx / C; + const int _c = idx % C; + const int _offset = _b * T * C + _c; + + F u = _u[_c]; + F w = _w[_c]; + const F *__restrict__ const k = _k + _offset; + const F *__restrict__ const v = _v + _offset; + F *__restrict__ const y = _y + _offset; + + // aa and bb are running sums divided by exp(pp) (to avoid overflow) + F aa = 0, bb = 0, pp = MIN_VALUE; + for (int i = 0; i < T; i++) { + const int ii = i * C; + const F kk = k[ii]; + const F vv = v[ii]; + + F ww = u + kk; + F p = max(pp, ww); + F e1 = exp(pp - p); + F e2 = exp(ww - p); + y[ii] = (e1 * aa + e2 * vv) / (e1 * bb + e2); + + ww = w + pp; + p = max(ww, kk); + e1 = exp(ww - p); + e2 = exp(kk - p); + aa = e1 * aa + e2 * vv; + bb = e1 * bb + e2; + pp = p; + } +} + +template +__global__ void kernel_backward(const int B, const int T, const int C, + const F *__restrict__ const _w, const F *__restrict__ const _u, const F *__restrict__ const _k, const F *__restrict__ const _v, + const F *__restrict__ const _y, const F *__restrict__ const _gy, + F *__restrict__ const _gw, F *__restrict__ const _gu, F *__restrict__ const _gk, F *__restrict__ const _gv) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + const int _b = idx / C; + const int _c = idx % C; + const int _offset = _b * T * C + _c; + + F u = _u[_c]; + F w = _w[_c]; + const F *__restrict__ const k = _k + _offset; + const F *__restrict__ const v = _v + _offset; + const F *__restrict__ const y = _y + _offset; + const F *__restrict__ const gy = _gy + _offset; + F *__restrict__ const gk = _gk + _offset; + F *__restrict__ const gv = _gv + _offset; + + F q[Tmax], r[Tmax]; + + F gw = 0, gu = 0, aa = 0, bb = 0, ga = 0, gb = 0, pp = MIN_VALUE; + for (int i = 0; i < T; i++) { + const int ii = i * C; + const F kk = k[ii]; + const F vv = v[ii]; + const F yy = y[ii]; + + F ww = u + kk; + F p = max(pp, ww); + F e1 = exp(pp - p); + F e2 = exp(ww - p); + const F qq = gy[ii] / (e1 * bb + e2); + gw += (ga - gb * yy) * e1 * qq; + gu += (vv - yy) * e2 * qq; + q[i] = qq; + r[i] = ww - p; + + ww = w + pp; + p = max(ww, kk); + e1 = exp(ww - p); + e2 = exp(kk - p); + ga = e1 * (aa + ga); + gb = e1 * (bb + gb); + aa = e1 * aa + e2 * vv; + bb = e1 * bb + e2; + pp = p; + } + const int _offsetBC = _b * C + _c; + _gw[_offsetBC] = gw * _w[_c]; // multiply by w because of w -> -exp(w) in python forward() + _gu[_offsetBC] = gu; + + aa = 0, bb = 0, pp = MIN_VALUE; + for (int i = T - 1; i >= 0; i--) { + const int ii = i * C; + const F kk = k[ii]; + const F vv = v[ii]; + const F yy = y[ii]; + const F qq = q[i]; + const F rr = r[i]; + + F e1 = qq * exp(rr); + F e2 = exp(kk + pp); + gk[ii] = e1 * (vv - yy) + e2 * (aa * vv + bb); + gv[ii] = e1 + e2 * aa; + + const F ww = w + pp; + const F www = rr - u - kk; + const F p = max(ww, www); + e1 = exp(ww - p); + e2 = qq * exp(www - p); + aa = e1 * aa + e2; + bb = e1 * bb - e2 * yy; + pp = p; + } +} + +void cuda_forward(int B, int T, int C, float *w, float *u, float *k, float *v, float *y) { + dim3 threadsPerBlock( min(C, 32) ); // requires --maxrregcount 60 for optimal performance + assert(B * C % threadsPerBlock.x == 0); + dim3 numBlocks(B * C / threadsPerBlock.x); + kernel_forward<<>>(B, T, C, w, u, k, v, y); +} + +void cuda_backward(int B, int T, int C, float *w, float *u, float *k, float *v, float *y, float *gy, float *gw, float *gu, float *gk, float *gv) { + dim3 threadsPerBlock( min(C, 32) ); // requires --maxrregcount 60 for optimal performance + assert(B * C % threadsPerBlock.x == 0); + dim3 numBlocks(B * C / threadsPerBlock.x); + kernel_backward<<>>(B, T, C, w, u, k, v, y, gy, gw, gu, gk, gv); +} diff --git a/funasr/models/rwkv_bat/cuda_decoder/wkv_op.cpp b/funasr/models/rwkv_bat/cuda_decoder/wkv_op.cpp new file mode 100644 index 0000000000000000000000000000000000000000..1024219999ffd7ced493a63802911e6c52c0507a --- /dev/null +++ b/funasr/models/rwkv_bat/cuda_decoder/wkv_op.cpp @@ -0,0 +1,37 @@ +/* + * Bsed on https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v4/cuda/wkv_op.cpp + Function signatures were modified based on https://github.com/huggingface/transformers/blob/main/src/transformers/kernels/rwkv/wkv_op.cpp + + */ + +#include + +void cuda_forward(int B, int T, int C, float *w, float *u, float *k, float *v, float *y); + +void cuda_backward(int B, int T, int C, float *w, float *u, float *k, float *v, float *y, float *gy, float *gw, float *gu, float *gk, float *gv); + +void forward(torch::Tensor &w, torch::Tensor &u, torch::Tensor &k, torch::Tensor &v, torch::Tensor &y) { + const int B = k.size(0); + const int T = k.size(1); + const int C = k.size(2); + + cuda_forward(B, T, C, w.data_ptr(), u.data_ptr(), k.data_ptr(), v.data_ptr(), y.data_ptr()); +} + +void backward(torch::Tensor &w, torch::Tensor &u, torch::Tensor &k, torch::Tensor &v, torch::Tensor &y, torch::Tensor &gy, torch::Tensor &gw, torch::Tensor &gu, torch::Tensor &gk, torch::Tensor &gv) { + const int B = k.size(0); + const int T = k.size(1); + const int C = k.size(2); + + cuda_backward(B, T, C, w.data_ptr(), u.data_ptr(), k.data_ptr(), v.data_ptr(), y.data_ptr(), gy.data_ptr(), gw.data_ptr(), gu.data_ptr(), gk.data_ptr(), gv.data_ptr()); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("forward", &forward, "wkv forward"); + m.def("backward", &backward, "wkv backward"); +} + +TORCH_LIBRARY(wkv_decoder, m) { + m.def("forward", forward); + m.def("backward", backward); +} diff --git a/funasr/models/rwkv_bat/cuda_encoder/wkv_cuda.cu b/funasr/models/rwkv_bat/cuda_encoder/wkv_cuda.cu new file mode 100644 index 0000000000000000000000000000000000000000..1dcbc7141dc69655534214f6df642731c8c1260f --- /dev/null +++ b/funasr/models/rwkv_bat/cuda_encoder/wkv_cuda.cu @@ -0,0 +1,135 @@ +// Copied from https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v4/cuda/wkv_cuda.cu + +#include +#include + +#define MIN_VALUE (-1e38) + +template +__global__ void kernel_forward(const int B, const int T, const int C, + const F *__restrict__ const _w, const F *__restrict__ const _u, const F *__restrict__ const _k, const F *__restrict__ const _v, + F *__restrict__ const _y) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + const int _b = idx / C; + const int _c = idx % C; + const int _offset = _b * T * C + _c; + + F u = _u[_c]; + F w = _w[_c]; + const F *__restrict__ const k = _k + _offset; + const F *__restrict__ const v = _v + _offset; + F *__restrict__ const y = _y + _offset; + + // aa and bb are running sums divided by exp(pp) (to avoid overflow) + F aa = 0, bb = 0, pp = MIN_VALUE; + for (int i = 0; i < T; i++) { + const int ii = i * C; + const F kk = k[ii]; + const F vv = v[ii]; + + F ww = u + kk; + F p = max(pp, ww); + F e1 = exp(pp - p); + F e2 = exp(ww - p); + y[ii] = (e1 * aa + e2 * vv) / (e1 * bb + e2); + + ww = w + pp; + p = max(ww, kk); + e1 = exp(ww - p); + e2 = exp(kk - p); + aa = e1 * aa + e2 * vv; + bb = e1 * bb + e2; + pp = p; + } +} + +template +__global__ void kernel_backward(const int B, const int T, const int C, + const F *__restrict__ const _w, const F *__restrict__ const _u, const F *__restrict__ const _k, const F *__restrict__ const _v, + const F *__restrict__ const _y, const F *__restrict__ const _gy, + F *__restrict__ const _gw, F *__restrict__ const _gu, F *__restrict__ const _gk, F *__restrict__ const _gv) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + const int _b = idx / C; + const int _c = idx % C; + const int _offset = _b * T * C + _c; + + F u = _u[_c]; + F w = _w[_c]; + const F *__restrict__ const k = _k + _offset; + const F *__restrict__ const v = _v + _offset; + const F *__restrict__ const y = _y + _offset; + const F *__restrict__ const gy = _gy + _offset; + F *__restrict__ const gk = _gk + _offset; + F *__restrict__ const gv = _gv + _offset; + + F q[Tmax], r[Tmax]; + + F gw = 0, gu = 0, aa = 0, bb = 0, ga = 0, gb = 0, pp = MIN_VALUE; + for (int i = 0; i < T; i++) { + const int ii = i * C; + const F kk = k[ii]; + const F vv = v[ii]; + const F yy = y[ii]; + + F ww = u + kk; + F p = max(pp, ww); + F e1 = exp(pp - p); + F e2 = exp(ww - p); + const F qq = gy[ii] / (e1 * bb + e2); + gw += (ga - gb * yy) * e1 * qq; + gu += (vv - yy) * e2 * qq; + q[i] = qq; + r[i] = ww - p; + + ww = w + pp; + p = max(ww, kk); + e1 = exp(ww - p); + e2 = exp(kk - p); + ga = e1 * (aa + ga); + gb = e1 * (bb + gb); + aa = e1 * aa + e2 * vv; + bb = e1 * bb + e2; + pp = p; + } + const int _offsetBC = _b * C + _c; + _gw[_offsetBC] = gw * _w[_c]; // multiply by w because of w -> -exp(w) in python forward() + _gu[_offsetBC] = gu; + + aa = 0, bb = 0, pp = MIN_VALUE; + for (int i = T - 1; i >= 0; i--) { + const int ii = i * C; + const F kk = k[ii]; + const F vv = v[ii]; + const F yy = y[ii]; + const F qq = q[i]; + const F rr = r[i]; + + F e1 = qq * exp(rr); + F e2 = exp(kk + pp); + gk[ii] = e1 * (vv - yy) + e2 * (aa * vv + bb); + gv[ii] = e1 + e2 * aa; + + const F ww = w + pp; + const F www = rr - u - kk; + const F p = max(ww, www); + e1 = exp(ww - p); + e2 = qq * exp(www - p); + aa = e1 * aa + e2; + bb = e1 * bb - e2 * yy; + pp = p; + } +} + +void cuda_forward(int B, int T, int C, float *w, float *u, float *k, float *v, float *y) { + dim3 threadsPerBlock( min(C, 32) ); // requires --maxrregcount 60 for optimal performance + assert(B * C % threadsPerBlock.x == 0); + dim3 numBlocks(B * C / threadsPerBlock.x); + kernel_forward<<>>(B, T, C, w, u, k, v, y); +} + +void cuda_backward(int B, int T, int C, float *w, float *u, float *k, float *v, float *y, float *gy, float *gw, float *gu, float *gk, float *gv) { + dim3 threadsPerBlock( min(C, 32) ); // requires --maxrregcount 60 for optimal performance + assert(B * C % threadsPerBlock.x == 0); + dim3 numBlocks(B * C / threadsPerBlock.x); + kernel_backward<<>>(B, T, C, w, u, k, v, y, gy, gw, gu, gk, gv); +} diff --git a/funasr/models/rwkv_bat/cuda_encoder/wkv_op.cpp b/funasr/models/rwkv_bat/cuda_encoder/wkv_op.cpp new file mode 100644 index 0000000000000000000000000000000000000000..16f364637eb12623c3ba9b982493794fd7f01a89 --- /dev/null +++ b/funasr/models/rwkv_bat/cuda_encoder/wkv_op.cpp @@ -0,0 +1,37 @@ +/* + * Bsed on https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v4/cuda/wkv_op.cpp + Function signatures were modified based on https://github.com/huggingface/transformers/blob/main/src/transformers/kernels/rwkv/wkv_op.cpp + + */ + +#include + +void cuda_forward(int B, int T, int C, float *w, float *u, float *k, float *v, float *y); + +void cuda_backward(int B, int T, int C, float *w, float *u, float *k, float *v, float *y, float *gy, float *gw, float *gu, float *gk, float *gv); + +void forward(torch::Tensor &w, torch::Tensor &u, torch::Tensor &k, torch::Tensor &v, torch::Tensor &y) { + const int B = k.size(0); + const int T = k.size(1); + const int C = k.size(2); + + cuda_forward(B, T, C, w.data_ptr(), u.data_ptr(), k.data_ptr(), v.data_ptr(), y.data_ptr()); +} + +void backward(torch::Tensor &w, torch::Tensor &u, torch::Tensor &k, torch::Tensor &v, torch::Tensor &y, torch::Tensor &gy, torch::Tensor &gw, torch::Tensor &gu, torch::Tensor &gk, torch::Tensor &gv) { + const int B = k.size(0); + const int T = k.size(1); + const int C = k.size(2); + + cuda_backward(B, T, C, w.data_ptr(), u.data_ptr(), k.data_ptr(), v.data_ptr(), y.data_ptr(), gy.data_ptr(), gw.data_ptr(), gu.data_ptr(), gk.data_ptr(), gv.data_ptr()); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("forward", &forward, "wkv forward"); + m.def("backward", &backward, "wkv backward"); +} + +TORCH_LIBRARY(wkv_encoder, m) { + m.def("forward", forward); + m.def("backward", backward); +} diff --git a/funasr/models/rwkv_bat/rwkv.py b/funasr/models/rwkv_bat/rwkv.py new file mode 100644 index 0000000000000000000000000000000000000000..c777cf2b87badb1abcf49a0c69d9442f667229ec --- /dev/null +++ b/funasr/models/rwkv_bat/rwkv.py @@ -0,0 +1,145 @@ +#!/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 typing import Dict, Optional, Tuple + +from funasr.models.transformer.layer_norm import LayerNorm +from funasr.models.rwkv_bat.rwkv_feed_forward import FeedForward +from funasr.models.rwkv_bat.rwkv_attention import EncoderSelfAttention, DecoderSelfAttention + + +class RWKV(torch.nn.Module): + """RWKV module. + + Args: + size: Input/Output size. + linear_size: Feed-forward hidden size. + attention_size: SelfAttention hidden size. + context_size: Context size for WKV computation. + block_id: Block index. + num_blocks: Number of blocks in the architecture. + normalization_class: Normalization layer class. + normalization_args: Normalization layer arguments. + att_dropout_rate: Dropout rate for the attention module. + ffn_dropout_rate: Dropout rate for the feed-forward module. + + """ + + def __init__( + self, + size: int, + linear_size: int, + attention_size: int, + context_size: int, + block_id: int, + num_blocks: int, + att_dropout_rate: float = 0.0, + ffn_dropout_rate: float = 0.0, + dropout_rate: float = 0.0, + ) -> None: + """Construct a RWKV object.""" + super().__init__() + + self.layer_norm_att = LayerNorm(size) + self.layer_norm_ffn = LayerNorm(size) + + self.att = EncoderSelfAttention( + size, attention_size, context_size, block_id, att_dropout_rate, num_blocks + ) + self.dropout_att = torch.nn.Dropout(p=dropout_rate) + + self.ffn = FeedForward(size, linear_size, block_id, ffn_dropout_rate, num_blocks) + self.dropout_ffn = torch.nn.Dropout(p=dropout_rate) + + def forward( + self, + x: torch.Tensor, + state: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Compute receptance weighted key value. + + Args: + x: RWKV input sequences. (B, L, size) + state: Decoder hidden states. [5 x (B, D_att/size, N)] + + Returns: + x: RWKV output sequences. (B, L, size) + x: Decoder hidden states. [5 x (B, D_att/size, N)] + + """ + att, state = self.att(self.layer_norm_att(x), state=state) + x = x + self.dropout_att(att) + ffn, state = self.ffn(self.layer_norm_ffn(x), state=state) + x = x + self.dropout_ffn(ffn) + return x, state + + +class RWKVDecoderLayer(torch.nn.Module): + """RWKV module. + + Args: + size: Input/Output size. + linear_size: Feed-forward hidden size. + attention_size: SelfAttention hidden size. + context_size: Context size for WKV computation. + block_id: Block index. + num_blocks: Number of blocks in the architecture. + normalization_class: Normalization layer class. + normalization_args: Normalization layer arguments. + att_dropout_rate: Dropout rate for the attention module. + ffn_dropout_rate: Dropout rate for the feed-forward module. + + """ + + def __init__( + self, + size: int, + linear_size: int, + attention_size: int, + context_size: int, + block_id: int, + num_blocks: int, + att_dropout_rate: float = 0.0, + ffn_dropout_rate: float = 0.0, + dropout_rate: float = 0.0, + ) -> None: + """Construct a RWKV object.""" + super().__init__() + + self.layer_norm_att = LayerNorm(size) + self.layer_norm_ffn = LayerNorm(size) + + self.att = DecoderSelfAttention( + size, attention_size, context_size, block_id, att_dropout_rate, num_blocks + ) + self.dropout_att = torch.nn.Dropout(p=dropout_rate) + + self.ffn = FeedForward(size, linear_size, block_id, ffn_dropout_rate, num_blocks) + self.dropout_ffn = torch.nn.Dropout(p=dropout_rate) + + def forward( + self, + x: torch.Tensor, + state: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Compute receptance weighted key value. + + Args: + x: RWKV input sequences. (B, L, size) + state: Decoder hidden states. [5 x (B, D_att/size, N)] + + Returns: + x: RWKV output sequences. (B, L, size) + x: Decoder hidden states. [5 x (B, D_att/size, N)] + + """ + att, state = self.att(self.layer_norm_att(x), state=state) + x = x + self.dropout_att(att) + + ffn, state = self.ffn(self.layer_norm_ffn(x), state=state) + x = x + self.dropout_ffn(ffn) + + return x, state diff --git a/funasr/models/rwkv_bat/rwkv_attention.py b/funasr/models/rwkv_bat/rwkv_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..59bf0ffbde1a0ff73e3ee38dad150a78f4544c03 --- /dev/null +++ b/funasr/models/rwkv_bat/rwkv_attention.py @@ -0,0 +1,608 @@ +#!/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 math +import torch +from pathlib import Path +from importlib.util import find_spec +from typing import List, Optional, Tuple, Union + + +wkv_kernel_encoder = None +wkv_kernel_decoder = None + + +class WKVLinearAttentionEncoder(torch.autograd.Function): + """WKVLinearAttention function definition.""" + + @staticmethod + def forward( + ctx, + time_decay: torch.Tensor, + time_first: torch.Tensor, + key: torch.Tensor, + value: torch.tensor, + ) -> torch.Tensor: + """WKVLinearAttention function forward pass. + + Args: + time_decay: Channel-wise time decay vector. (D_att) + time_first: Channel-wise time first vector. (D_att) + key: Key tensor. (B, U, D_att) + value: Value tensor. (B, U, D_att) + + Returns: + out: Weighted Key-Value tensor. (B, U, D_att) + + """ + batch, length, dim = key.size() + + assert length <= wkv_kernel_encoder.context_size, ( + f"Cannot process key of length {length} while context_size " + f"is ({wkv_kernel_encoder.context_size}). Limit should be increased." + ) + + assert batch * dim % min(dim, 32) == 0, ( + f"batch size ({batch}) by dimension ({dim}) should be a multiple of " f"{min(dim, 32)}" + ) + + ctx.input_dtype = key.dtype + + time_decay = -torch.exp(time_decay.float().contiguous()) + time_first = time_first.float().contiguous() + + key = key.float().contiguous() + value = value.float().contiguous() + + out = torch.empty_like(key, memory_format=torch.contiguous_format) + + wkv_kernel_encoder.forward(time_decay, time_first, key, value, out) + ctx.save_for_backward(time_decay, time_first, key, value, out) + + return out + + @staticmethod + def backward( + ctx, grad_output: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """WKVLinearAttention function backward pass. + + Args: + grad_output: Output gradient. (B, U, D_att) + + Returns: + grad_time_decay: Gradient for channel-wise time decay vector. (D_att) + grad_time_first: Gradient for channel-wise time first vector. (D_att) + grad_key: Gradient for key tensor. (B, U, D_att) + grad_value: Gradient for value tensor. (B, U, D_att) + + """ + time_decay, time_first, key, value, output = ctx.saved_tensors + grad_dtype = ctx.input_dtype + + batch, _, dim = key.size() + + grad_time_decay = torch.empty( + (batch, dim), + memory_format=torch.contiguous_format, + dtype=time_decay.dtype, + device=time_decay.device, + ) + + grad_time_first = torch.empty( + (batch, dim), + memory_format=torch.contiguous_format, + dtype=time_decay.dtype, + device=time_decay.device, + ) + + grad_key = torch.empty_like(key, memory_format=torch.contiguous_format) + grad_value = torch.empty_like(value, memory_format=torch.contiguous_format) + + wkv_kernel_encoder.backward( + time_decay, + time_first, + key, + value, + output, + grad_output.contiguous(), + grad_time_decay, + grad_time_first, + grad_key, + grad_value, + ) + + grad_time_decay = torch.sum(grad_time_decay, dim=0) + grad_time_first = torch.sum(grad_time_first, dim=0) + + return ( + grad_time_decay, + grad_time_first, + grad_key, + grad_value, + ) + + +class WKVLinearAttentionDecoder(torch.autograd.Function): + """WKVLinearAttention function definition.""" + + @staticmethod + def forward( + ctx, + time_decay: torch.Tensor, + time_first: torch.Tensor, + key: torch.Tensor, + value: torch.tensor, + ) -> torch.Tensor: + """WKVLinearAttention function forward pass. + + Args: + time_decay: Channel-wise time decay vector. (D_att) + time_first: Channel-wise time first vector. (D_att) + key: Key tensor. (B, U, D_att) + value: Value tensor. (B, U, D_att) + + Returns: + out: Weighted Key-Value tensor. (B, U, D_att) + + """ + batch, length, dim = key.size() + + assert length <= wkv_kernel_decoder.context_size, ( + f"Cannot process key of length {length} while context_size " + f"is ({wkv_kernel.context_size}). Limit should be increased." + ) + + assert batch * dim % min(dim, 32) == 0, ( + f"batch size ({batch}) by dimension ({dim}) should be a multiple of " f"{min(dim, 32)}" + ) + + ctx.input_dtype = key.dtype + + time_decay = -torch.exp(time_decay.float().contiguous()) + time_first = time_first.float().contiguous() + + key = key.float().contiguous() + value = value.float().contiguous() + + out = torch.empty_like(key, memory_format=torch.contiguous_format) + + wkv_kernel_decoder.forward(time_decay, time_first, key, value, out) + ctx.save_for_backward(time_decay, time_first, key, value, out) + + return out + + @staticmethod + def backward( + ctx, grad_output: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """WKVLinearAttention function backward pass. + + Args: + grad_output: Output gradient. (B, U, D_att) + + Returns: + grad_time_decay: Gradient for channel-wise time decay vector. (D_att) + grad_time_first: Gradient for channel-wise time first vector. (D_att) + grad_key: Gradient for key tensor. (B, U, D_att) + grad_value: Gradient for value tensor. (B, U, D_att) + + """ + time_decay, time_first, key, value, output = ctx.saved_tensors + grad_dtype = ctx.input_dtype + + batch, _, dim = key.size() + + grad_time_decay = torch.empty( + (batch, dim), + memory_format=torch.contiguous_format, + dtype=time_decay.dtype, + device=time_decay.device, + ) + + grad_time_first = torch.empty( + (batch, dim), + memory_format=torch.contiguous_format, + dtype=time_decay.dtype, + device=time_decay.device, + ) + + grad_key = torch.empty_like(key, memory_format=torch.contiguous_format) + grad_value = torch.empty_like(value, memory_format=torch.contiguous_format) + + wkv_kernel_decoder.backward( + time_decay, + time_first, + key, + value, + output, + grad_output.contiguous(), + grad_time_decay, + grad_time_first, + grad_key, + grad_value, + ) + + grad_time_decay = torch.sum(grad_time_decay, dim=0) + grad_time_first = torch.sum(grad_time_first, dim=0) + + return ( + grad_time_decay, + grad_time_first, + grad_key, + grad_value, + ) + + +def load_encoder_wkv_kernel(context_size: int) -> None: + """Load WKV CUDA kernel. + + Args: + context_size: Context size. + + """ + from torch.utils.cpp_extension import load + + global wkv_kernel_encoder + + if wkv_kernel_encoder is not None and wkv_kernel_encoder.context_size == context_size: + return + + if find_spec("ninja") is None: + raise ImportError( + "Ninja package was not found. WKV kernel module can't be loaded " + "for training. Please, 'pip install ninja' in your environment." + ) + + if not torch.cuda.is_available(): + raise ImportError( + "CUDA is currently a requirement for WKV kernel loading. " + "Please set your devices properly and launch again." + ) + + kernel_folder = Path(__file__).resolve().parent / "cuda_encoder" + kernel_files = [kernel_folder / f for f in ["wkv_op.cpp", "wkv_cuda.cu"]] + + kernel_cflags = [ + "-res-usage", + "--maxrregcount 60", + "--use_fast_math", + "-O3", + "-Xptxas -O3", + f"-DTmax={context_size}", + ] + wkv_kernel_encoder = load( + name=f"encoder_wkv_{context_size}", + sources=kernel_files, + verbose=True, + extra_cuda_cflags=kernel_cflags, + ) + wkv_kernel_encoder.context_size = context_size + + +def load_decoder_wkv_kernel(context_size: int) -> None: + """Load WKV CUDA kernel. + + Args: + context_size: Context size. + + """ + from torch.utils.cpp_extension import load + + global wkv_kernel_decoder + + if wkv_kernel_decoder is not None and wkv_kernel_decoder.context_size == context_size: + return + + if find_spec("ninja") is None: + raise ImportError( + "Ninja package was not found. WKV kernel module can't be loaded " + "for training. Please, 'pip install ninja' in your environment." + ) + + if not torch.cuda.is_available(): + raise ImportError( + "CUDA is currently a requirement for WKV kernel loading. " + "Please set your devices properly and launch again." + ) + + kernel_folder = Path(__file__).resolve().parent / "cuda_decoder" + kernel_files = [kernel_folder / f for f in ["wkv_op.cpp", "wkv_cuda.cu"]] + + kernel_cflags = [ + "-res-usage", + "--maxrregcount 60", + "--use_fast_math", + "-O3", + "-Xptxas -O3", + f"-DTmax={context_size}", + ] + wkv_kernel_decoder = load( + name=f"decoder_wkv_{context_size}", + sources=kernel_files, + verbose=True, + extra_cuda_cflags=kernel_cflags, + ) + wkv_kernel_decoder.context_size = context_size + + +class SelfAttention(torch.nn.Module): + """SelfAttention module definition. + + Args: + size: Input/Output size. + attention_size: Attention hidden size. + context_size: Context size for WKV kernel. + block_id: Block index. + num_blocks: Number of blocks in the architecture. + + """ + + def __init__( + self, + size: int, + attention_size: int, + block_id: int, + dropout_rate: float, + num_blocks: int, + ) -> None: + """Construct a SelfAttention object.""" + super().__init__() + self.time_shift = torch.nn.ZeroPad2d((0, 0, 1, -1)) + + self.time_decay = torch.nn.Parameter(torch.empty(attention_size)) + self.time_first = torch.nn.Parameter(torch.empty(attention_size)) + + self.time_mix_key = torch.nn.Parameter(torch.empty(1, 1, size)) + self.time_mix_value = torch.nn.Parameter(torch.empty(1, 1, size)) + self.time_mix_receptance = torch.nn.Parameter(torch.empty(1, 1, size)) + + self.proj_key = torch.nn.Linear(size, attention_size, bias=True) + self.proj_value = torch.nn.Linear(size, attention_size, bias=True) + self.proj_receptance = torch.nn.Linear(size, attention_size, bias=True) + + self.proj_output = torch.nn.Linear(attention_size, size, bias=True) + + self.block_id = block_id + + self.reset_parameters(size, attention_size, block_id, num_blocks) + self.dropout = torch.nn.Dropout(p=dropout_rate) + + def reset_parameters( + self, size: int, attention_size: int, block_id: int, num_blocks: int + ) -> None: + """Reset module parameters. + + Args: + size: Block size. + attention_size: Attention hidden size. + block_id: Block index. + num_blocks: Number of blocks in the architecture. + + """ + ratio_0_to_1 = block_id / (num_blocks - 1) + ratio_1_to_almost0 = 1.0 - (block_id / num_blocks) + + time_weight = torch.ones(1, 1, size) + + for i in range(size): + time_weight[0, 0, i] = i / size + + decay_speed = [ + -5 + 8 * (h / (attention_size - 1)) ** (0.7 + 1.3 * ratio_0_to_1) + for h in range(attention_size) + ] + decay_speed = torch.tensor( + decay_speed, dtype=self.time_decay.dtype, device=self.time_decay.device + ) + + zigzag = ( + torch.tensor( + [(i + 1) % 3 - 1 for i in range(attention_size)], + dtype=self.time_first.dtype, + device=self.time_first.device, + ) + * 0.5 + ) + + with torch.no_grad(): + self.time_decay.data = decay_speed + self.time_first.data = torch.ones_like(self.time_first * math.log(0.3) + zigzag) + + self.time_mix_key.data = torch.pow(time_weight, ratio_1_to_almost0) + self.time_mix_value.data = ( + torch.pow(time_weight, ratio_1_to_almost0) + 0.3 * ratio_0_to_1 + ) + self.time_mix_receptance.data = torch.pow(time_weight, 0.5 * ratio_1_to_almost0) + + @torch.no_grad() + def wkv_linear_attention( + self, + time_decay: torch.Tensor, + time_first: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + state: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], + ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + """Compute WKV with state (i.e.: for inference). + + Args: + time_decay: Channel-wise time decay vector. (D_att) + time_first: Channel-wise time first vector. (D_att) + key: Key tensor. (B, 1, D_att) + value: Value tensor. (B, 1, D_att) + state: Decoder hidden states. [3 x (B, D_att)] + + Returns: + output: Weighted Key-Value. (B, 1, D_att) + state: Decoder hidden states. [3 x (B, 1, D_att)] + + """ + num_state, den_state, max_state = state + time_decay = -torch.exp(time_decay) + max_for_output = torch.maximum(max_state, (time_first + key)) + + e1 = torch.exp(max_state - max_for_output) + e2 = torch.exp((time_first + key) - max_for_output) + + numerator = e1 * num_state + e2 * value + denominator = e1 * den_state + e2 + + max_for_state = torch.maximum(key, (max_state + time_decay)) + + e1 = torch.exp((max_state + time_decay) - max_for_state) + e2 = torch.exp(key - max_for_state) + + wkv = numerator / denominator + + state = [e1 * num_state + e2 * value, e1 * den_state + e2, max_for_state] + + return wkv, state + + +class DecoderSelfAttention(SelfAttention): + """SelfAttention module definition. + + Args: + size: Input/Output size. + attention_size: Attention hidden size. + context_size: Context size for WKV kernel. + block_id: Block index. + num_blocks: Number of blocks in the architecture. + + """ + + def __init__( + self, + size: int, + attention_size: int, + context_size: int, + block_id: int, + dropout_rate: float, + num_blocks: int, + ) -> None: + """Construct a SelfAttention object.""" + super().__init__(size, attention_size, block_id, dropout_rate, num_blocks) + # load_decoder_wkv_kernel(context_size) + + def forward( + self, + x: torch.Tensor, + state: Optional[List[torch.Tensor]] = None, + ) -> Tuple[torch.Tensor, Optional[List[torch.Tensor]]]: + """Compute time mixing. + + Args: + x: SelfAttention input sequences. (B, U, size) + state: Decoder hidden states. [5 x (B, 1, D_att, N)] + + Returns: + x: SelfAttention output sequences. (B, U, size) + + """ + shifted_x = self.time_shift(x) if state is None else state[1][..., self.block_id] + + key = x * self.time_mix_key + shifted_x * (1 - self.time_mix_key) + value = x * self.time_mix_value + shifted_x * (1 - self.time_mix_value) + receptance = x * self.time_mix_receptance + shifted_x * (1 - self.time_mix_receptance) + + key = self.proj_key(key) + value = self.proj_value(value) + receptance = torch.sigmoid(self.proj_receptance(receptance)) + + if state is not None: + state[1][..., self.block_id] = x + + wkv, att_state = self.wkv_linear_attention( + self.time_decay, + self.time_first, + key, + value, + tuple(s[..., self.block_id] for s in state[2:]), + ) + + state[2][..., self.block_id] = att_state[0] + state[3][..., self.block_id] = att_state[1] + state[4][..., self.block_id] = att_state[2] + else: + wkv = WKVLinearAttentionDecoder.apply(self.time_decay, self.time_first, key, value) + + wkv = self.dropout(wkv) + x = self.proj_output(receptance * wkv) + + return x, state + + +class EncoderSelfAttention(SelfAttention): + """SelfAttention module definition. + + Args: + size: Input/Output size. + attention_size: Attention hidden size. + context_size: Context size for WKV kernel. + block_id: Block index. + num_blocks: Number of blocks in the architecture. + + """ + + def __init__( + self, + size: int, + attention_size: int, + context_size: int, + block_id: int, + dropout_rate: float, + num_blocks: int, + ) -> None: + """Construct a SelfAttention object.""" + super().__init__(size, attention_size, block_id, dropout_rate, num_blocks) + # load_encoder_wkv_kernel(context_size) + + def forward( + self, + x: torch.Tensor, + state: Optional[List[torch.Tensor]] = None, + ) -> Tuple[torch.Tensor, Optional[List[torch.Tensor]]]: + """Compute time mixing. + + Args: + x: SelfAttention input sequences. (B, U, size) + state: Decoder hidden states. [5 x (B, 1, D_att, N)] + + Returns: + x: SelfAttention output sequences. (B, U, size) + + """ + shifted_x = self.time_shift(x) if state is None else state[1][..., self.block_id] + + key = x * self.time_mix_key + shifted_x * (1 - self.time_mix_key) + value = x * self.time_mix_value + shifted_x * (1 - self.time_mix_value) + receptance = x * self.time_mix_receptance + shifted_x * (1 - self.time_mix_receptance) + + key = self.proj_key(key) + value = self.proj_value(value) + receptance = torch.sigmoid(self.proj_receptance(receptance)) + + if state is not None: + state[1][..., self.block_id] = x + + wkv, att_state = self.wkv_linear_attention( + self.time_decay, + self.time_first, + key, + value, + tuple(s[..., self.block_id] for s in state[2:]), + ) + + state[2][..., self.block_id] = att_state[0] + state[3][..., self.block_id] = att_state[1] + state[4][..., self.block_id] = att_state[2] + else: + wkv = WKVLinearAttentionEncoder.apply(self.time_decay, self.time_first, key, value) + + wkv = self.dropout(wkv) + x = self.proj_output(receptance * wkv) + + return x, state diff --git a/funasr/models/rwkv_bat/rwkv_encoder.py b/funasr/models/rwkv_bat/rwkv_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..a27088d69f6114b8c63856c98b33df7fd0b905a0 --- /dev/null +++ b/funasr/models/rwkv_bat/rwkv_encoder.py @@ -0,0 +1,158 @@ +#!/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 typing import Dict, List, Optional, Tuple + +from funasr.register import tables +from funasr.models.rwkv_bat.rwkv import RWKV +from funasr.models.transformer.layer_norm import LayerNorm +from funasr.models.transformer.utils.nets_utils import make_source_mask +from funasr.models.rwkv_bat.rwkv_subsampling import RWKVConvInput + + +@tables.register("encoder_classes", "RWKVEncoder") +class RWKVEncoder(torch.nn.Module): + """RWKV encoder module. + + Based on https://arxiv.org/pdf/2305.13048.pdf. + + Args: + vocab_size: Vocabulary size. + output_size: Input/Output size. + context_size: Context size for WKV computation. + linear_size: FeedForward hidden size. + attention_size: SelfAttention hidden size. + normalization_type: Normalization layer type. + normalization_args: Normalization layer arguments. + num_blocks: Number of RWKV blocks. + embed_dropout_rate: Dropout rate for embedding layer. + att_dropout_rate: Dropout rate for the attention module. + ffn_dropout_rate: Dropout rate for the feed-forward module. + """ + + def __init__( + self, + input_size: int, + output_size: int = 512, + context_size: int = 1024, + linear_size: Optional[int] = None, + attention_size: Optional[int] = None, + num_blocks: int = 4, + att_dropout_rate: float = 0.0, + ffn_dropout_rate: float = 0.0, + dropout_rate: float = 0.0, + subsampling_factor: int = 4, + time_reduction_factor: int = 1, + kernel: int = 3, + **kwargs, + ) -> None: + """Construct a RWKVEncoder object.""" + super().__init__() + + self.embed = RWKVConvInput( + input_size, + [output_size // 4, output_size // 2, output_size], + subsampling_factor, + conv_kernel_size=kernel, + output_size=output_size, + ) + + self.subsampling_factor = subsampling_factor + + linear_size = output_size * 4 if linear_size is None else linear_size + attention_size = output_size if attention_size is None else attention_size + + self.rwkv_blocks = torch.nn.ModuleList( + [ + RWKV( + output_size, + linear_size, + attention_size, + context_size, + block_id, + num_blocks, + att_dropout_rate=att_dropout_rate, + ffn_dropout_rate=ffn_dropout_rate, + dropout_rate=dropout_rate, + ) + for block_id in range(num_blocks) + ] + ) + + self.embed_norm = LayerNorm(output_size) + self.final_norm = LayerNorm(output_size) + + self._output_size = output_size + self.context_size = context_size + + self.num_blocks = num_blocks + self.time_reduction_factor = time_reduction_factor + + def output_size(self) -> int: + return self._output_size + + def forward(self, x: torch.Tensor, x_len) -> torch.Tensor: + """Encode source label sequences. + + Args: + x: Encoder input sequences. (B, L) + + Returns: + out: Encoder output sequences. (B, U, D) + + """ + _, length, _ = x.size() + + assert ( + length <= self.context_size * self.subsampling_factor + ), "Context size is too short for current length: %d versus %d" % ( + length, + self.context_size * self.subsampling_factor, + ) + mask = make_source_mask(x_len).to(x.device) + x, mask = self.embed(x, mask, None) + x = self.embed_norm(x) + olens = mask.eq(0).sum(1) + + if self.training: + for block in self.rwkv_blocks: + x, _ = block(x) + else: + x = self.rwkv_infer(x) + + x = self.final_norm(x) + + if self.time_reduction_factor > 1: + x = x[:, :: self.time_reduction_factor, :] + olens = torch.floor_divide(olens - 1, self.time_reduction_factor) + 1 + + return x, olens, None + + def rwkv_infer(self, xs_pad): + + batch_size = xs_pad.shape[0] + + hidden_sizes = [self._output_size for i in range(5)] + + state = [ + torch.zeros( + (batch_size, 1, hidden_sizes[i], self.num_blocks), + dtype=torch.float32, + device=xs_pad.device, + ) + for i in range(5) + ] + + state[4] -= 1e-30 + + xs_out = [] + for t in range(xs_pad.shape[1]): + x_t = xs_pad[:, t, :] + for idx, block in enumerate(self.rwkv_blocks): + x_t, state = block(x_t, state=state) + xs_out.append(x_t) + xs_out = torch.cat(xs_out, dim=1) + return xs_out diff --git a/funasr/models/rwkv_bat/rwkv_feed_forward.py b/funasr/models/rwkv_bat/rwkv_feed_forward.py new file mode 100644 index 0000000000000000000000000000000000000000..a8056bf7019f962906074055e3d809b4b971da00 --- /dev/null +++ b/funasr/models/rwkv_bat/rwkv_feed_forward.py @@ -0,0 +1,89 @@ +#!/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 typing import List, Optional, Tuple + + +class FeedForward(torch.nn.Module): + """FeedForward module definition. + + Args: + size: Input/Output size. + hidden_size: Hidden size. + block_id: Block index. + num_blocks: Number of blocks in the architecture. + + """ + + def __init__( + self, size: int, hidden_size: int, block_id: int, dropout_rate: float, num_blocks: int + ) -> None: + """Construct a FeedForward object.""" + super().__init__() + + self.time_shift = torch.nn.ZeroPad2d((0, 0, 1, -1)) + + self.time_mix_key = torch.nn.Parameter(torch.empty(1, 1, size)) + self.time_mix_receptance = torch.nn.Parameter(torch.empty(1, 1, size)) + + self.proj_key = torch.nn.Linear(size, hidden_size, bias=True) + self.proj_value = torch.nn.Linear(hidden_size, size, bias=True) + self.proj_receptance = torch.nn.Linear(size, size, bias=True) + + self.block_id = block_id + + self.reset_parameters(size, block_id, num_blocks) + self.dropout = torch.nn.Dropout(p=dropout_rate) + + def reset_parameters(self, size: int, block_id: int, num_blocks: int) -> None: + """Reset module parameters. + + Args: + size: Block size. + block_id: Block index. + num_blocks: Number of blocks in the architecture. + + """ + ratio_1_to_almost0 = 1.0 - (block_id / num_blocks) + + time_weight = torch.ones(1, 1, size) + + for i in range(size): + time_weight[0, 0, i] = i / size + + with torch.no_grad(): + self.time_mix_key.data = torch.pow(time_weight, ratio_1_to_almost0) + self.time_mix_receptance.data = torch.pow(time_weight, ratio_1_to_almost0) + + def forward( + self, x: torch.Tensor, state: Optional[List[torch.Tensor]] = None + ) -> Tuple[torch.Tensor, Optional[List[torch.Tensor]]]: + """Compute channel mixing. + + Args: + x: FeedForward input sequences. (B, U, size) + state: Decoder hidden state. [5 x (B, 1, size, N)] + + Returns: + x: FeedForward output sequences. (B, U, size) + state: Decoder hidden state. [5 x (B, 1, size, N)] + + """ + shifted_x = self.time_shift(x) if state is None else state[0][..., self.block_id] + + key = x * self.time_mix_key + shifted_x * (1 - self.time_mix_key) + receptance = x * self.time_mix_receptance + shifted_x * (1 - self.time_mix_receptance) + + key = torch.square(torch.relu(self.proj_key(key))) + value = self.proj_value(self.dropout(key)) + receptance = torch.sigmoid(self.proj_receptance(receptance)) + + if state is not None: + state[0][..., self.block_id] = x + + x = receptance * value + + return x, state diff --git a/funasr/models/rwkv_bat/rwkv_subsampling.py b/funasr/models/rwkv_bat/rwkv_subsampling.py new file mode 100644 index 0000000000000000000000000000000000000000..5108ae68db28dede990611d1f60dfcd1930eca73 --- /dev/null +++ b/funasr/models/rwkv_bat/rwkv_subsampling.py @@ -0,0 +1,250 @@ +#!/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 math +import torch +from typing import Optional, Tuple, Union +from funasr.models.transformer.utils.nets_utils import pad_to_len + + +class TooShortUttError(Exception): + """Raised when the utt is too short for subsampling. + + Args: + message (str): Message for error catch + actual_size (int): the short size that cannot pass the subsampling + limit (int): the limit size for subsampling + + """ + + def __init__(self, message, actual_size, limit): + """Construct a TooShortUttError for error handler.""" + super().__init__(message) + self.actual_size = actual_size + self.limit = limit + + +def check_short_utt(ins, size): + """Check if the utterance is too short for subsampling.""" + if isinstance(ins, Conv2dSubsampling2) and size < 3: + return True, 3 + if isinstance(ins, Conv2dSubsampling) and size < 7: + return True, 7 + if isinstance(ins, Conv2dSubsampling6) and size < 11: + return True, 11 + if isinstance(ins, Conv2dSubsampling8) and size < 15: + return True, 15 + return False, -1 + + +class RWKVConvInput(torch.nn.Module): + """Streaming ConvInput module definition. + Args: + input_size: Input size. + conv_size: Convolution size. + subsampling_factor: Subsampling factor. + output_size: Block output dimension. + """ + + def __init__( + self, + input_size: int, + conv_size: Union[int, Tuple], + subsampling_factor: int = 4, + conv_kernel_size: int = 3, + output_size: Optional[int] = None, + ) -> None: + """Construct a ConvInput object.""" + super().__init__() + if subsampling_factor == 1: + conv_size1, conv_size2, conv_size3 = conv_size + + self.conv = torch.nn.Sequential( + torch.nn.Conv2d( + 1, conv_size1, conv_kernel_size, stride=1, padding=(conv_kernel_size - 1) // 2 + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size1, + conv_size1, + conv_kernel_size, + stride=[1, 2], + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size1, + conv_size2, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size2, + conv_size2, + conv_kernel_size, + stride=[1, 2], + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size2, + conv_size3, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size3, + conv_size3, + conv_kernel_size, + stride=[1, 2], + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + ) + + output_proj = conv_size3 * ((input_size // 2) // 2) + + self.subsampling_factor = 1 + + self.stride_1 = 1 + + self.create_new_mask = self.create_new_vgg_mask + + else: + conv_size1, conv_size2, conv_size3 = conv_size + + kernel_1 = int(subsampling_factor / 2) + + self.conv = torch.nn.Sequential( + torch.nn.Conv2d( + 1, conv_size1, conv_kernel_size, stride=1, padding=(conv_kernel_size - 1) // 2 + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size1, + conv_size1, + conv_kernel_size, + stride=[kernel_1, 2], + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size1, + conv_size2, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size2, + conv_size2, + conv_kernel_size, + stride=[2, 2], + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size2, + conv_size3, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size3, + conv_size3, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + ) + + output_proj = conv_size3 * ((input_size // 2) // 2) + + self.subsampling_factor = subsampling_factor + + self.create_new_mask = self.create_new_vgg_mask + + self.stride_1 = kernel_1 + + self.min_frame_length = 7 + + if output_size is not None: + self.output = torch.nn.Linear(output_proj, output_size) + self.output_size = output_size + else: + self.output = None + self.output_size = output_proj + + def forward( + self, x: torch.Tensor, mask: Optional[torch.Tensor], chunk_size: Optional[torch.Tensor] + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Encode input sequences. + Args: + x: ConvInput input sequences. (B, T, D_feats) + mask: Mask of input sequences. (B, 1, T) + Returns: + x: ConvInput output sequences. (B, sub(T), D_out) + mask: Mask of output sequences. (B, 1, sub(T)) + """ + if mask is not None: + mask = self.create_new_mask(mask) + olens = max(mask.eq(0).sum(1)) + + b, t, f = x.size() + x = x.unsqueeze(1) # (b. 1. t. f) + + if chunk_size is not None: + max_input_length = int( + chunk_size + * self.subsampling_factor + * (math.ceil(float(t) / (chunk_size * self.subsampling_factor))) + ) + x = map(lambda inputs: pad_to_len(inputs, max_input_length, 1), x) + x = list(x) + x = torch.stack(x, dim=0) + N_chunks = max_input_length // (chunk_size * self.subsampling_factor) + x = x.view(b * N_chunks, 1, chunk_size * self.subsampling_factor, f) + + x = self.conv(x) + + _, c, _, f = x.size() + if chunk_size is not None: + x = x.transpose(1, 2).contiguous().view(b, -1, c * f)[:, :olens, :] + else: + x = x.transpose(1, 2).contiguous().view(b, -1, c * f) + + if self.output is not None: + x = self.output(x) + + return x, mask[:, :olens][:, : x.size(1)] + + def create_new_vgg_mask(self, mask: torch.Tensor) -> torch.Tensor: + """Create a new mask for VGG output sequences. + Args: + mask: Mask of input sequences. (B, T) + Returns: + mask: Mask of output sequences. (B, sub(T)) + """ + if self.subsampling_factor > 1: + return mask[:, ::2][:, :: self.stride_1] + else: + return mask + + def get_size_before_subsampling(self, size: int) -> int: + """Return the original size before subsampling for a given size. + Args: + size: Number of frames after subsampling. + Returns: + : Number of frames before subsampling. + """ + return size * self.subsampling_factor diff --git a/funasr/models/rwkv_bat/template.yaml b/funasr/models/rwkv_bat/template.yaml new file mode 100644 index 0000000000000000000000000000000000000000..dad84bf38baa9e6b96f2aa62036c22cbd24fa893 --- /dev/null +++ b/funasr/models/rwkv_bat/template.yaml @@ -0,0 +1,64 @@ +# network architecture +model: Transducer +model_conf: + auxiliary_ctc_weight: 0.0 + +# encoder +encoder: RWKVEncoder +encoder_conf: + kernel: 3 + subsampling_factor: 4 + output_size: 512 + num_blocks: 18 + time_reduction_factor: 2 + att_dropout_rate: 0.1 + ffn_dropout_rate: 0.1 + dropout_rate: 0.1 + +# decoder (prediction network) +decoder: rnnt_decoder +decoder_conf: + embed_size: 512 + hidden_size: 512 + embed_dropout_rate: 0.1 + dropout_rate: 0.1 + use_embed_mask: false + +# joint network +joint_network: joint_network +joint_network_conf: + joint_space_size: 512 + +frontend: WavFrontend +frontend_conf: + fs: 16000 + window: hamming + n_mels: 80 + frame_length: 25 + frame_shift: 10 + lfr_m: 1 + lfr_n: 1 + upsacle_samples: true + +specaug: SpecAugLFR +specaug_conf: + apply_time_warp: false + time_warp_window: 5 + time_warp_mode: bicubic + apply_freq_mask: true + freq_mask_width_range: + - 0 + - 30 + lfr_rate: 6 + num_freq_mask: 1 + apply_time_mask: true + time_mask_width_range: + - 0 + - 12 + num_time_mask: 1 + +tokenizer: CharTokenizer +tokenizer_conf: + unk_symbol: + split_with_space: true + \ No newline at end of file diff --git a/funasr/models/sa_asr/__init__.py b/funasr/models/sa_asr/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/models/sa_asr/attention.py b/funasr/models/sa_asr/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..74e687a521eb5e8eadab5bec733a9cea05269a0f --- /dev/null +++ b/funasr/models/sa_asr/attention.py @@ -0,0 +1,50 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + + +"""Multi-Head Attention layer definition.""" + +import math + +import numpy +import torch +from torch import nn +from typing import Optional, Tuple + +import torch.nn.functional as F +from funasr.models.transformer.utils.nets_utils import make_pad_mask +import funasr.models.lora.layers as lora + + +class CosineDistanceAttention(nn.Module): + """Compute Cosine Distance between spk decoder output and speaker profile + Args: + profile_path: speaker profile file path (.npy file) + """ + + def __init__(self): + super().__init__() + self.softmax = nn.Softmax(dim=-1) + + def forward(self, spk_decoder_out, profile, profile_lens=None): + """ + Args: + spk_decoder_out(torch.Tensor):(B, L, D) + spk_profiles(torch.Tensor):(B, N, D) + """ + x = spk_decoder_out.unsqueeze(2) # (B, L, 1, D) + if profile_lens is not None: + + mask = (make_pad_mask(profile_lens)[:, None, :]).to(profile.device) + min_value = float(numpy.finfo(torch.tensor(0, dtype=x.dtype).numpy().dtype).min) + weights_not_softmax = F.cosine_similarity(x, profile.unsqueeze(1), dim=-1).masked_fill( + mask, min_value + ) + weights = self.softmax(weights_not_softmax).masked_fill(mask, 0.0) # (B, L, N) + else: + x = x[:, -1:, :, :] + weights_not_softmax = F.cosine_similarity(x, profile.unsqueeze(1).to(x.device), dim=-1) + weights = self.softmax(weights_not_softmax) # (B, 1, N) + spk_embedding = torch.matmul(weights, profile.to(weights.device)) # (B, L, D) + + return spk_embedding, weights diff --git a/funasr/models/sa_asr/beam_search_sa_asr.py b/funasr/models/sa_asr/beam_search_sa_asr.py new file mode 100644 index 0000000000000000000000000000000000000000..d953b4c6a98d5b669ebedf146e53245e3d26b908 --- /dev/null +++ b/funasr/models/sa_asr/beam_search_sa_asr.py @@ -0,0 +1,540 @@ +"""Beam search module.""" + +from itertools import chain +import logging +from typing import Any +from typing import Dict +from typing import List +from typing import NamedTuple +from typing import Tuple +from typing import Union + +import torch + +from funasr.metrics import end_detect +from funasr.models.transformer.scorers.scorer_interface import PartialScorerInterface +from funasr.models.transformer.scorers.scorer_interface import ScorerInterface +from funasr.models.decoder.abs_decoder import AbsDecoder + + +class Hypothesis(NamedTuple): + """Hypothesis data type.""" + + yseq: torch.Tensor + spk_weigths: List + score: Union[float, torch.Tensor] = 0 + scores: Dict[str, Union[float, torch.Tensor]] = dict() + states: Dict[str, Any] = dict() + + def asdict(self) -> dict: + """Convert data to JSON-friendly dict.""" + return self._replace( + yseq=self.yseq.tolist(), + score=float(self.score), + scores={k: float(v) for k, v in self.scores.items()}, + )._asdict() + + +class BeamSearch(torch.nn.Module): + """Beam search implementation.""" + + def __init__( + self, + scorers: Dict[str, ScorerInterface], + weights: Dict[str, float], + beam_size: int, + vocab_size: int, + sos: int, + eos: int, + token_list: List[str] = None, + pre_beam_ratio: float = 1.5, + pre_beam_score_key: str = None, + ): + """Initialize beam search. + + Args: + scorers (dict[str, ScorerInterface]): Dict of decoder modules + e.g., Decoder, CTCPrefixScorer, LM + The scorer will be ignored if it is `None` + weights (dict[str, float]): Dict of weights for each scorers + The scorer will be ignored if its weight is 0 + beam_size (int): The number of hypotheses kept during search + vocab_size (int): The number of vocabulary + sos (int): Start of sequence id + eos (int): End of sequence id + token_list (list[str]): List of tokens for debug log + pre_beam_score_key (str): key of scores to perform pre-beam search + pre_beam_ratio (float): beam size in the pre-beam search + will be `int(pre_beam_ratio * beam_size)` + + """ + super().__init__() + # set scorers + self.weights = weights + self.scorers = dict() + self.full_scorers = dict() + self.part_scorers = dict() + # this module dict is required for recursive cast + # `self.to(device, dtype)` in `recog.py` + self.nn_dict = torch.nn.ModuleDict() + for k, v in scorers.items(): + w = weights.get(k, 0) + if w == 0 or v is None: + continue + assert isinstance( + v, ScorerInterface + ), f"{k} ({type(v)}) does not implement ScorerInterface" + self.scorers[k] = v + if isinstance(v, PartialScorerInterface): + self.part_scorers[k] = v + else: + self.full_scorers[k] = v + if isinstance(v, torch.nn.Module): + self.nn_dict[k] = v + + # set configurations + self.sos = sos + self.eos = eos + self.token_list = token_list + self.pre_beam_size = int(pre_beam_ratio * beam_size) + self.beam_size = beam_size + self.n_vocab = vocab_size + if ( + pre_beam_score_key is not None + and pre_beam_score_key != "full" + and pre_beam_score_key not in self.full_scorers + ): + raise KeyError(f"{pre_beam_score_key} is not found in {self.full_scorers}") + self.pre_beam_score_key = pre_beam_score_key + self.do_pre_beam = ( + self.pre_beam_score_key is not None + and self.pre_beam_size < self.n_vocab + and len(self.part_scorers) > 0 + ) + + def init_hyp(self, x: torch.Tensor) -> List[Hypothesis]: + """Get an initial hypothesis data. + + Args: + x (torch.Tensor): The encoder output feature + + Returns: + Hypothesis: The initial hypothesis. + + """ + init_states = dict() + init_scores = dict() + for k, d in self.scorers.items(): + init_states[k] = d.init_state(x) + init_scores[k] = 0.0 + return [ + Hypothesis( + score=0.0, + scores=init_scores, + states=init_states, + yseq=torch.tensor([self.sos], device=x.device), + spk_weigths=[], + ) + ] + + @staticmethod + def append_token(xs: torch.Tensor, x: int) -> torch.Tensor: + """Append new token to prefix tokens. + + Args: + xs (torch.Tensor): The prefix token + x (int): The new token to append + + Returns: + torch.Tensor: New tensor contains: xs + [x] with xs.dtype and xs.device + + """ + x = torch.tensor([x], dtype=xs.dtype, device=xs.device) + return torch.cat((xs, x)) + + def score_full( + self, + hyp: Hypothesis, + asr_enc: torch.Tensor, + spk_enc: torch.Tensor, + profile: torch.Tensor, + ) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: + """Score new hypothesis by `self.full_scorers`. + + Args: + hyp (Hypothesis): Hypothesis with prefix tokens to score + x (torch.Tensor): Corresponding input feature + + Returns: + Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: Tuple of + score dict of `hyp` that has string keys of `self.full_scorers` + and tensor score values of shape: `(self.n_vocab,)`, + and state dict that has string keys + and state values of `self.full_scorers` + + """ + scores = dict() + states = dict() + for k, d in self.full_scorers.items(): + if isinstance(d, AbsDecoder): + scores[k], spk_weigths, states[k] = d.score( + hyp.yseq, hyp.states[k], asr_enc, spk_enc, profile + ) + else: + scores[k], states[k] = d.score(hyp.yseq, hyp.states[k], asr_enc) + return scores, spk_weigths, states + + def score_partial( + self, + hyp: Hypothesis, + ids: torch.Tensor, + asr_enc: torch.Tensor, + spk_enc: torch.Tensor, + profile: torch.Tensor, + ) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: + """Score new hypothesis by `self.part_scorers`. + + Args: + hyp (Hypothesis): Hypothesis with prefix tokens to score + ids (torch.Tensor): 1D tensor of new partial tokens to score + x (torch.Tensor): Corresponding input feature + + Returns: + Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: Tuple of + score dict of `hyp` that has string keys of `self.part_scorers` + and tensor score values of shape: `(len(ids),)`, + and state dict that has string keys + and state values of `self.part_scorers` + + """ + scores = dict() + states = dict() + for k, d in self.part_scorers.items(): + if isinstance(d, AbsDecoder): + scores[k], states[k] = d.score_partial( + hyp.yseq, ids, hyp.states[k], asr_enc, spk_enc, profile + ) + else: + scores[k], states[k] = d.score_partial(hyp.yseq, ids, hyp.states[k], asr_enc) + return scores, states + + def beam( + self, weighted_scores: torch.Tensor, ids: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute topk full token ids and partial token ids. + + Args: + weighted_scores (torch.Tensor): The weighted sum scores for each tokens. + Its shape is `(self.n_vocab,)`. + ids (torch.Tensor): The partial token ids to compute topk + + Returns: + Tuple[torch.Tensor, torch.Tensor]: + The topk full token ids and partial token ids. + Their shapes are `(self.beam_size,)` + + """ + # no pre beam performed + if weighted_scores.size(0) == ids.size(0): + top_ids = weighted_scores.topk(self.beam_size)[1] + return top_ids, top_ids + + # mask pruned in pre-beam not to select in topk + tmp = weighted_scores[ids] + weighted_scores[:] = -float("inf") + weighted_scores[ids] = tmp + top_ids = weighted_scores.topk(self.beam_size)[1] + local_ids = weighted_scores[ids].topk(self.beam_size)[1] + return top_ids, local_ids + + @staticmethod + def merge_scores( + prev_scores: Dict[str, float], + next_full_scores: Dict[str, torch.Tensor], + full_idx: int, + next_part_scores: Dict[str, torch.Tensor], + part_idx: int, + ) -> Dict[str, torch.Tensor]: + """Merge scores for new hypothesis. + + Args: + prev_scores (Dict[str, float]): + The previous hypothesis scores by `self.scorers` + next_full_scores (Dict[str, torch.Tensor]): scores by `self.full_scorers` + full_idx (int): The next token id for `next_full_scores` + next_part_scores (Dict[str, torch.Tensor]): + scores of partial tokens by `self.part_scorers` + part_idx (int): The new token id for `next_part_scores` + + Returns: + Dict[str, torch.Tensor]: The new score dict. + Its keys are names of `self.full_scorers` and `self.part_scorers`. + Its values are scalar tensors by the scorers. + + """ + new_scores = dict() + for k, v in next_full_scores.items(): + new_scores[k] = prev_scores[k] + v[full_idx] + for k, v in next_part_scores.items(): + new_scores[k] = prev_scores[k] + v[part_idx] + return new_scores + + def merge_states(self, states: Any, part_states: Any, part_idx: int) -> Any: + """Merge states for new hypothesis. + + Args: + states: states of `self.full_scorers` + part_states: states of `self.part_scorers` + part_idx (int): The new token id for `part_scores` + + Returns: + Dict[str, torch.Tensor]: The new score dict. + Its keys are names of `self.full_scorers` and `self.part_scorers`. + Its values are states of the scorers. + + """ + new_states = dict() + for k, v in states.items(): + new_states[k] = v + for k, d in self.part_scorers.items(): + new_states[k] = d.select_state(part_states[k], part_idx) + return new_states + + def search( + self, + running_hyps: List[Hypothesis], + asr_enc: torch.Tensor, + spk_enc: torch.Tensor, + profile: torch.Tensor, + ) -> List[Hypothesis]: + """Search new tokens for running hypotheses and encoded speech x. + + Args: + running_hyps (List[Hypothesis]): Running hypotheses on beam + x (torch.Tensor): Encoded speech feature (T, D) + + Returns: + List[Hypotheses]: Best sorted hypotheses + + """ + # import ipdb;ipdb.set_trace() + best_hyps = [] + part_ids = torch.arange(self.n_vocab, device=asr_enc.device) # no pre-beam + for hyp in running_hyps: + # scoring + weighted_scores = torch.zeros(self.n_vocab, dtype=asr_enc.dtype, device=asr_enc.device) + scores, spk_weigths, states = self.score_full(hyp, asr_enc, spk_enc, profile) + for k in self.full_scorers: + weighted_scores += self.weights[k] * scores[k] + # partial scoring + if self.do_pre_beam: + pre_beam_scores = ( + weighted_scores + if self.pre_beam_score_key == "full" + else scores[self.pre_beam_score_key] + ) + part_ids = torch.topk(pre_beam_scores, self.pre_beam_size)[1] + part_scores, part_states = self.score_partial(hyp, part_ids, asr_enc, spk_enc, profile) + for k in self.part_scorers: + weighted_scores[part_ids] += self.weights[k] * part_scores[k] + # add previous hyp score + weighted_scores += hyp.score + + # update hyps + for j, part_j in zip(*self.beam(weighted_scores, part_ids)): + # will be (2 x beam at most) + best_hyps.append( + Hypothesis( + score=weighted_scores[j], + yseq=self.append_token(hyp.yseq, j), + scores=self.merge_scores(hyp.scores, scores, j, part_scores, part_j), + states=self.merge_states(states, part_states, part_j), + spk_weigths=hyp.spk_weigths + [spk_weigths], + ) + ) + + # sort and prune 2 x beam -> beam + best_hyps = sorted(best_hyps, key=lambda x: x.score, reverse=True)[ + : min(len(best_hyps), self.beam_size) + ] + return best_hyps + + def forward( + self, + asr_enc: torch.Tensor, + spk_enc: torch.Tensor, + profile: torch.Tensor, + maxlenratio: float = 0.0, + minlenratio: float = 0.0, + ) -> List[Hypothesis]: + """Perform beam search. + + Args: + x (torch.Tensor): Encoded speech feature (T, D) + maxlenratio (float): Input length ratio to obtain max output length. + If maxlenratio=0.0 (default), it uses a end-detect function + to automatically find maximum hypothesis lengths + minlenratio (float): Input length ratio to obtain min output length. + + Returns: + list[Hypothesis]: N-best decoding results + + """ + # import ipdb;ipdb.set_trace() + # set length bounds + if maxlenratio == 0: + maxlen = asr_enc.shape[0] + else: + maxlen = max(1, int(maxlenratio * asr_enc.size(0))) + minlen = int(minlenratio * asr_enc.size(0)) + logging.info("decoder input length: " + str(asr_enc.shape[0])) + logging.info("max output length: " + str(maxlen)) + logging.info("min output length: " + str(minlen)) + + # main loop of prefix search + running_hyps = self.init_hyp(asr_enc) + ended_hyps = [] + for i in range(maxlen): + logging.debug("position " + str(i)) + best = self.search(running_hyps, asr_enc, spk_enc, profile) + # import pdb;pdb.set_trace() + # post process of one iteration + running_hyps = self.post_process(i, maxlen, maxlenratio, best, ended_hyps) + # end detection + if maxlenratio == 0.0 and end_detect([h.asdict() for h in ended_hyps], i): + logging.info(f"end detected at {i}") + break + if len(running_hyps) == 0: + logging.info("no hypothesis. Finish decoding.") + break + else: + logging.debug(f"remained hypotheses: {len(running_hyps)}") + + nbest_hyps = sorted(ended_hyps, key=lambda x: x.score, reverse=True) + # check the number of hypotheses reaching to eos + if len(nbest_hyps) == 0: + logging.warning( + "there is no N-best results, perform recognition " "again with smaller minlenratio." + ) + return ( + [] + if minlenratio < 0.1 + else self.forward( + asr_enc, spk_enc, profile, maxlenratio, max(0.0, minlenratio - 0.1) + ) + ) + + # report the best result + best = nbest_hyps[0] + for k, v in best.scores.items(): + logging.info(f"{v:6.2f} * {self.weights[k]:3} = {v * self.weights[k]:6.2f} for {k}") + logging.info(f"total log probability: {best.score:.2f}") + logging.info(f"normalized log probability: {best.score / len(best.yseq):.2f}") + logging.info(f"total number of ended hypotheses: {len(nbest_hyps)}") + if self.token_list is not None: + logging.info( + "best hypo: " + "".join([self.token_list[x] for x in best.yseq[1:-1]]) + "\n" + ) + return nbest_hyps + + def post_process( + self, + i: int, + maxlen: int, + maxlenratio: float, + running_hyps: List[Hypothesis], + ended_hyps: List[Hypothesis], + ) -> List[Hypothesis]: + """Perform post-processing of beam search iterations. + + Args: + i (int): The length of hypothesis tokens. + maxlen (int): The maximum length of tokens in beam search. + maxlenratio (int): The maximum length ratio in beam search. + running_hyps (List[Hypothesis]): The running hypotheses in beam search. + ended_hyps (List[Hypothesis]): The ended hypotheses in beam search. + + Returns: + List[Hypothesis]: The new running hypotheses. + + """ + logging.debug(f"the number of running hypotheses: {len(running_hyps)}") + if self.token_list is not None: + logging.debug( + "best hypo: " + "".join([self.token_list[x] for x in running_hyps[0].yseq[1:]]) + ) + # add eos in the final loop to avoid that there are no ended hyps + if i == maxlen - 1: + logging.info("adding in the last position in the loop") + running_hyps = [ + h._replace(yseq=self.append_token(h.yseq, self.eos)) for h in running_hyps + ] + + # add ended hypotheses to a final list, and removed them from current hypotheses + # (this will be a problem, number of hyps < beam) + remained_hyps = [] + for hyp in running_hyps: + if hyp.yseq[-1] == self.eos: + # e.g., Word LM needs to add final score + for k, d in chain(self.full_scorers.items(), self.part_scorers.items()): + s = d.final_score(hyp.states[k]) + hyp.scores[k] += s + hyp = hyp._replace(score=hyp.score + self.weights[k] * s) + ended_hyps.append(hyp) + else: + remained_hyps.append(hyp) + return remained_hyps + + +def beam_search( + x: torch.Tensor, + sos: int, + eos: int, + beam_size: int, + vocab_size: int, + scorers: Dict[str, ScorerInterface], + weights: Dict[str, float], + token_list: List[str] = None, + maxlenratio: float = 0.0, + minlenratio: float = 0.0, + pre_beam_ratio: float = 1.5, + pre_beam_score_key: str = "full", +) -> list: + """Perform beam search with scorers. + + Args: + x (torch.Tensor): Encoded speech feature (T, D) + sos (int): Start of sequence id + eos (int): End of sequence id + beam_size (int): The number of hypotheses kept during search + vocab_size (int): The number of vocabulary + scorers (dict[str, ScorerInterface]): Dict of decoder modules + e.g., Decoder, CTCPrefixScorer, LM + The scorer will be ignored if it is `None` + weights (dict[str, float]): Dict of weights for each scorers + The scorer will be ignored if its weight is 0 + token_list (list[str]): List of tokens for debug log + maxlenratio (float): Input length ratio to obtain max output length. + If maxlenratio=0.0 (default), it uses a end-detect function + to automatically find maximum hypothesis lengths + minlenratio (float): Input length ratio to obtain min output length. + pre_beam_score_key (str): key of scores to perform pre-beam search + pre_beam_ratio (float): beam size in the pre-beam search + will be `int(pre_beam_ratio * beam_size)` + + Returns: + list: N-best decoding results + + """ + ret = BeamSearch( + scorers, + weights, + beam_size=beam_size, + vocab_size=vocab_size, + pre_beam_ratio=pre_beam_ratio, + pre_beam_score_key=pre_beam_score_key, + sos=sos, + eos=eos, + token_list=token_list, + ).forward(x=x, maxlenratio=maxlenratio, minlenratio=minlenratio) + return [h.asdict() for h in ret] diff --git a/funasr/models/sa_asr/e2e_sa_asr.py b/funasr/models/sa_asr/e2e_sa_asr.py new file mode 100644 index 0000000000000000000000000000000000000000..f4827b25784418c40533cbd39fcae75325e36704 --- /dev/null +++ b/funasr/models/sa_asr/e2e_sa_asr.py @@ -0,0 +1,508 @@ +# Copyright ESPnet (https://github.com/espnet/espnet). All Rights Reserved. +# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +import logging +from contextlib import contextmanager +from distutils.version import LooseVersion +from typing import Dict +from typing import List +from typing import Optional +from typing import Tuple +from typing import Union + +import torch +import torch.nn.functional as F + +from funasr.layers.abs_normalize import AbsNormalize +from funasr.losses.label_smoothing_loss import LabelSmoothingLoss, NllLoss # noqa: H301 +from funasr.models.ctc import CTC +from funasr.models.decoder.abs_decoder import AbsDecoder +from funasr.models.encoder.abs_encoder import AbsEncoder +from funasr.frontends.abs_frontend import AbsFrontend +from funasr.models.postencoder.abs_postencoder import AbsPostEncoder +from funasr.models.preencoder.abs_preencoder import AbsPreEncoder +from funasr.models.specaug.abs_specaug import AbsSpecAug +from funasr.models.transformer.utils.add_sos_eos import add_sos_eos +from funasr.metrics import ErrorCalculator +from funasr.metrics.compute_acc import th_accuracy +from funasr.train_utils.device_funcs import force_gatherable +from funasr.models.base_model import FunASRModel + +if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): + from torch.cuda.amp import autocast +else: + # Nothing to do if torch<1.6.0 + @contextmanager + def autocast(enabled=True): + yield + + +class SAASRModel(FunASRModel): + """CTC-attention hybrid Encoder-Decoder model""" + + def __init__( + self, + vocab_size: int, + max_spk_num: int, + token_list: Union[Tuple[str, ...], List[str]], + frontend: Optional[AbsFrontend], + specaug: Optional[AbsSpecAug], + normalize: Optional[AbsNormalize], + asr_encoder: AbsEncoder, + spk_encoder: torch.nn.Module, + decoder: AbsDecoder, + ctc: CTC, + spk_weight: float = 0.5, + ctc_weight: float = 0.5, + interctc_weight: float = 0.0, + ignore_id: int = -1, + lsm_weight: float = 0.0, + length_normalized_loss: bool = False, + report_cer: bool = True, + report_wer: bool = True, + sym_space: str = "", + sym_blank: str = "", + extract_feats_in_collect_stats: bool = True, + ): + assert 0.0 <= ctc_weight <= 1.0, ctc_weight + assert 0.0 <= interctc_weight < 1.0, interctc_weight + + super().__init__() + # note that eos is the same as sos (equivalent ID) + self.blank_id = 0 + self.sos = 1 + self.eos = 2 + self.vocab_size = vocab_size + self.max_spk_num = max_spk_num + self.ignore_id = ignore_id + self.spk_weight = spk_weight + self.ctc_weight = ctc_weight + self.interctc_weight = interctc_weight + self.token_list = token_list.copy() + + self.frontend = frontend + self.specaug = specaug + self.normalize = normalize + self.asr_encoder = asr_encoder + self.spk_encoder = spk_encoder + + if not hasattr(self.asr_encoder, "interctc_use_conditioning"): + self.asr_encoder.interctc_use_conditioning = False + if self.asr_encoder.interctc_use_conditioning: + self.asr_encoder.conditioning_layer = torch.nn.Linear( + vocab_size, self.asr_encoder.output_size() + ) + + self.error_calculator = None + + # we set self.decoder = None in the CTC mode since + # self.decoder parameters were never used and PyTorch complained + # and threw an Exception in the multi-GPU experiment. + # thanks Jeff Farris for pointing out the issue. + if ctc_weight == 1.0: + self.decoder = None + else: + self.decoder = decoder + + self.criterion_att = LabelSmoothingLoss( + size=vocab_size, + padding_idx=ignore_id, + smoothing=lsm_weight, + normalize_length=length_normalized_loss, + ) + + self.criterion_spk = NllLoss( + size=max_spk_num, + padding_idx=ignore_id, + normalize_length=length_normalized_loss, + ) + + if report_cer or report_wer: + self.error_calculator = ErrorCalculator( + token_list, sym_space, sym_blank, report_cer, report_wer + ) + + if ctc_weight == 0.0: + self.ctc = None + else: + self.ctc = ctc + + self.extract_feats_in_collect_stats = extract_feats_in_collect_stats + + def forward( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + text: torch.Tensor, + text_lengths: torch.Tensor, + profile: torch.Tensor, + profile_lengths: torch.Tensor, + text_id: torch.Tensor, + text_id_lengths: torch.Tensor, + ) -> 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,) + profile: (Batch, Length, Dim) + profile_lengths: (Batch,) + """ + assert text_lengths.dim() == 1, text_lengths.shape + # Check that batch_size is unified + assert ( + speech.shape[0] == speech_lengths.shape[0] == text.shape[0] == text_lengths.shape[0] + ), (speech.shape, speech_lengths.shape, text.shape, text_lengths.shape) + batch_size = speech.shape[0] + + # for data-parallel + text = text[:, : text_lengths.max()] + + # 1. Encoder + asr_encoder_out, encoder_out_lens, spk_encoder_out = self.encode(speech, speech_lengths) + intermediate_outs = None + if isinstance(asr_encoder_out, tuple): + intermediate_outs = asr_encoder_out[1] + asr_encoder_out = asr_encoder_out[0] + + loss_att, loss_spk, acc_att, acc_spk, cer_att, wer_att = None, None, None, None, None, None + loss_ctc, cer_ctc = None, None + stats = dict() + + # 1. CTC branch + if self.ctc_weight != 0.0: + loss_ctc, cer_ctc = self._calc_ctc_loss( + asr_encoder_out, encoder_out_lens, text, text_lengths + ) + + # Intermediate CTC (optional) + loss_interctc = 0.0 + if self.interctc_weight != 0.0 and intermediate_outs is not None: + for layer_idx, intermediate_out in intermediate_outs: + # we assume intermediate_out has the same length & padding + # as those of encoder_out + loss_ic, cer_ic = self._calc_ctc_loss( + intermediate_out, encoder_out_lens, text, text_lengths + ) + loss_interctc = loss_interctc + loss_ic + + # Collect Intermedaite CTC stats + stats["loss_interctc_layer{}".format(layer_idx)] = ( + loss_ic.detach() if loss_ic is not None else None + ) + stats["cer_interctc_layer{}".format(layer_idx)] = cer_ic + + loss_interctc = loss_interctc / len(intermediate_outs) + + # calculate whole encoder loss + loss_ctc = (1 - self.interctc_weight) * loss_ctc + self.interctc_weight * loss_interctc + + # 2b. Attention decoder branch + if self.ctc_weight != 1.0: + loss_att, loss_spk, acc_att, acc_spk, cer_att, wer_att = self._calc_att_loss( + asr_encoder_out, + spk_encoder_out, + encoder_out_lens, + text, + text_lengths, + profile, + profile_lengths, + text_id, + text_id_lengths, + ) + + # 3. CTC-Att loss definition + if self.ctc_weight == 0.0: + loss_asr = loss_att + elif self.ctc_weight == 1.0: + loss_asr = loss_ctc + else: + loss_asr = self.ctc_weight * loss_ctc + (1 - self.ctc_weight) * loss_att + + if self.spk_weight == 0.0: + loss = loss_asr + else: + loss = self.spk_weight * loss_spk + (1 - self.spk_weight) * loss_asr + + stats = dict( + loss=loss.detach(), + loss_asr=loss_asr.detach(), + loss_att=loss_att.detach() if loss_att is not None else None, + loss_ctc=loss_ctc.detach() if loss_ctc is not None else None, + loss_spk=loss_spk.detach() if loss_spk is not None else None, + acc=acc_att, + acc_spk=acc_spk, + cer=cer_att, + wer=wer_att, + cer_ctc=cer_ctc, + ) + + # force_gatherable: to-device and to-tensor if scalar for DataParallel + 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 + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Frontend + Encoder. Note that this method is used by asr_inference.py + + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + """ + with autocast(False): + # 1. Extract feats + feats, feats_lengths = self._extract_feats(speech, speech_lengths) + + # 2. Data augmentation + feats_raw = feats.clone() + if self.specaug is not None and self.training: + feats, feats_lengths = self.specaug(feats, feats_lengths) + + # 3. Normalization for feature: e.g. Global-CMVN, Utterance-CMVN + if self.normalize is not None: + feats, feats_lengths = self.normalize(feats, feats_lengths) + + # 4. Forward encoder + # feats: (Batch, Length, Dim) + # -> encoder_out: (Batch, Length2, Dim2) + if self.asr_encoder.interctc_use_conditioning: + encoder_out, encoder_out_lens, _ = self.asr_encoder(feats, feats_lengths, ctc=self.ctc) + else: + encoder_out, encoder_out_lens, _ = self.asr_encoder(feats, feats_lengths) + intermediate_outs = None + if isinstance(encoder_out, tuple): + intermediate_outs = encoder_out[1] + encoder_out = encoder_out[0] + + encoder_out_spk_ori = self.spk_encoder(feats_raw, feats_lengths)[0] + # import ipdb;ipdb.set_trace() + if encoder_out_spk_ori.size(1) != encoder_out.size(1): + encoder_out_spk = F.interpolate( + encoder_out_spk_ori.transpose(-2, -1), size=(encoder_out.size(1)), mode="nearest" + ).transpose(-2, -1) + else: + encoder_out_spk = encoder_out_spk_ori + + assert encoder_out.size(0) == speech.size(0), ( + encoder_out.size(), + speech.size(0), + ) + assert encoder_out.size(1) <= encoder_out_lens.max(), ( + encoder_out.size(), + encoder_out_lens.max(), + ) + assert encoder_out_spk.size(0) == speech.size(0), ( + encoder_out_spk.size(), + speech.size(0), + ) + + if intermediate_outs is not None: + return (encoder_out, intermediate_outs), encoder_out_lens, encoder_out_spk + + return encoder_out, encoder_out_lens, encoder_out_spk + + def _extract_feats( + self, speech: torch.Tensor, speech_lengths: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + assert speech_lengths.dim() == 1, speech_lengths.shape + + # for data-parallel + speech = speech[:, : speech_lengths.max()] + + if self.frontend is not None: + # Frontend + # e.g. STFT and Feature extract + # data_loader may send time-domain signal in this case + # speech (Batch, NSamples) -> feats: (Batch, NFrames, Dim) + feats, feats_lengths = self.frontend(speech, speech_lengths) + else: + # No frontend and no feature extract + feats, feats_lengths = speech, speech_lengths + return feats, feats_lengths + + 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, + asr_encoder_out: torch.Tensor, + spk_encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, + profile: torch.Tensor, + profile_lens: torch.Tensor, + text_id: torch.Tensor, + text_id_lengths: 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, weights_no_pad, _ = self.decoder( + asr_encoder_out, + spk_encoder_out, + encoder_out_lens, + ys_in_pad, + ys_in_lens, + profile, + profile_lens, + ) + + spk_num_no_pad = weights_no_pad.size(-1) + pad = (0, self.max_spk_num - spk_num_no_pad) + weights = F.pad(weights_no_pad, pad, mode="constant", value=0) + + # pre_id=weights.argmax(-1) + # pre_text=decoder_out.argmax(-1) + # id_mask=(pre_id==text_id).to(dtype=text_id.dtype) + # pre_text_mask=pre_text*id_mask+1-id_mask #相同的地方不变，不同的地方设为1() + # padding_mask= ys_out_pad != self.ignore_id + # numerator = torch.sum(pre_text_mask.masked_select(padding_mask) == ys_out_pad.masked_select(padding_mask)) + # denominator = torch.sum(padding_mask) + # sd_acc = float(numerator) / float(denominator) + + # 2. Compute attention loss + loss_att = self.criterion_att(decoder_out, ys_out_pad) + loss_spk = self.criterion_spk(torch.log(weights), text_id) + + acc_spk = th_accuracy( + weights.view(-1, self.max_spk_num), + text_id, + ignore_label=self.ignore_id, + ) + 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, loss_spk, acc_att, acc_spk, cer_att, wer_att + + def _calc_ctc_loss( + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, + ): + # Calc CTC loss + loss_ctc = self.ctc(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens) + + # Calc CER using CTC + cer_ctc = None + if not self.training and self.error_calculator is not None: + ys_hat = self.ctc.argmax(encoder_out).data + cer_ctc = self.error_calculator(ys_hat.cpu(), ys_pad.cpu(), is_ctc=True) + return loss_ctc, cer_ctc diff --git a/funasr/models/sa_asr/transformer_decoder.py b/funasr/models/sa_asr/transformer_decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..be33fa527c5b1a8490769a45837f028fcc9bea44 --- /dev/null +++ b/funasr/models/sa_asr/transformer_decoder.py @@ -0,0 +1,731 @@ +# Copyright 2019 Shigeki Karita +# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +"""Decoder definition.""" +from typing import Any +from typing import List +from typing import Sequence +from typing import Tuple + +import torch +from torch import nn + + +from funasr.models.transformer.attention import MultiHeadedAttention +from funasr.models.sa_asr.attention import CosineDistanceAttention +from funasr.models.transformer.utils.dynamic_conv import DynamicConvolution +from funasr.models.transformer.utils.dynamic_conv2d import DynamicConvolution2D +from funasr.models.transformer.embedding import PositionalEncoding +from funasr.models.transformer.layer_norm import LayerNorm +from funasr.models.transformer.utils.lightconv import LightweightConvolution +from funasr.models.transformer.utils.lightconv2d import LightweightConvolution2D +from funasr.models.transformer.utils.mask import subsequent_mask +from funasr.models.transformer.utils.nets_utils import make_pad_mask +from funasr.models.transformer.positionwise_feed_forward import ( + PositionwiseFeedForward, # noqa: H301 +) +from funasr.models.transformer.utils.repeat import repeat +from funasr.models.transformer.scorers.scorer_interface import BatchScorerInterface + +from funasr.register import tables + + +class DecoderLayer(nn.Module): + """Single decoder layer module. + + Args: + size (int): Input dimension. + self_attn (torch.nn.Module): Self-attention module instance. + `MultiHeadedAttention` instance can be used as the argument. + src_attn (torch.nn.Module): Self-attention module instance. + `MultiHeadedAttention` instance can be used as the argument. + feed_forward (torch.nn.Module): Feed-forward module instance. + `PositionwiseFeedForward`, `MultiLayeredConv1d`, or `Conv1dLinear` instance + can be used as the argument. + dropout_rate (float): Dropout rate. + normalize_before (bool): Whether to use layer_norm before the first block. + concat_after (bool): Whether to concat attention layer's input and output. + if True, additional linear will be applied. + i.e. x -> x + linear(concat(x, att(x))) + if False, no additional linear will be applied. i.e. x -> x + att(x) + + + """ + + def __init__( + self, + size, + self_attn, + src_attn, + feed_forward, + dropout_rate, + normalize_before=True, + concat_after=False, + ): + """Construct an DecoderLayer object.""" + super(DecoderLayer, self).__init__() + self.size = size + self.self_attn = self_attn + self.src_attn = src_attn + self.feed_forward = feed_forward + self.norm1 = LayerNorm(size) + self.norm2 = LayerNorm(size) + self.norm3 = LayerNorm(size) + self.dropout = nn.Dropout(dropout_rate) + self.normalize_before = normalize_before + self.concat_after = concat_after + if self.concat_after: + self.concat_linear1 = nn.Linear(size + size, size) + self.concat_linear2 = nn.Linear(size + size, size) + + def forward(self, tgt, tgt_mask, memory, memory_mask, cache=None): + """Compute decoded features. + + Args: + tgt (torch.Tensor): Input tensor (#batch, maxlen_out, size). + tgt_mask (torch.Tensor): Mask for input tensor (#batch, maxlen_out). + memory (torch.Tensor): Encoded memory, float32 (#batch, maxlen_in, size). + memory_mask (torch.Tensor): Encoded memory mask (#batch, maxlen_in). + cache (List[torch.Tensor]): List of cached tensors. + Each tensor shape should be (#batch, maxlen_out - 1, size). + + Returns: + torch.Tensor: Output tensor(#batch, maxlen_out, size). + torch.Tensor: Mask for output tensor (#batch, maxlen_out). + torch.Tensor: Encoded memory (#batch, maxlen_in, size). + torch.Tensor: Encoded memory mask (#batch, maxlen_in). + + """ + residual = tgt + if self.normalize_before: + tgt = self.norm1(tgt) + + if cache is None: + tgt_q = tgt + tgt_q_mask = tgt_mask + else: + # compute only the last frame query keeping dim: max_time_out -> 1 + assert cache.shape == ( + tgt.shape[0], + tgt.shape[1] - 1, + self.size, + ), f"{cache.shape} == {(tgt.shape[0], tgt.shape[1] - 1, self.size)}" + tgt_q = tgt[:, -1:, :] + residual = residual[:, -1:, :] + tgt_q_mask = None + if tgt_mask is not None: + tgt_q_mask = tgt_mask[:, -1:, :] + + if self.concat_after: + tgt_concat = torch.cat((tgt_q, self.self_attn(tgt_q, tgt, tgt, tgt_q_mask)), dim=-1) + x = residual + self.concat_linear1(tgt_concat) + else: + x = residual + self.dropout(self.self_attn(tgt_q, tgt, tgt, tgt_q_mask)) + if not self.normalize_before: + x = self.norm1(x) + + residual = x + if self.normalize_before: + x = self.norm2(x) + if self.concat_after: + x_concat = torch.cat((x, self.src_attn(x, memory, memory, memory_mask)), dim=-1) + x = residual + self.concat_linear2(x_concat) + else: + x = residual + self.dropout(self.src_attn(x, memory, memory, memory_mask)) + if not self.normalize_before: + x = self.norm2(x) + + residual = x + if self.normalize_before: + x = self.norm3(x) + x = residual + self.dropout(self.feed_forward(x)) + if not self.normalize_before: + x = self.norm3(x) + + if cache is not None: + x = torch.cat([cache, x], dim=1) + + return x, tgt_mask, memory, memory_mask + + +class BaseTransformerDecoder(nn.Module, BatchScorerInterface): + """Base class of Transfomer decoder module. + + Args: + vocab_size: output dim + encoder_output_size: dimension of attention + attention_heads: the number of heads of multi head attention + linear_units: the number of units of position-wise feed forward + num_blocks: the number of decoder blocks + dropout_rate: dropout rate + self_attention_dropout_rate: dropout rate for attention + input_layer: input layer type + use_output_layer: whether to use output layer + pos_enc_class: PositionalEncoding or ScaledPositionalEncoding + normalize_before: whether to use layer_norm before the first block + concat_after: whether to concat attention layer's input and output + if True, additional linear will be applied. + i.e. x -> x + linear(concat(x, att(x))) + if False, no additional linear will be applied. + i.e. x -> x + att(x) + """ + + def __init__( + self, + vocab_size: int, + encoder_output_size: int, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + ): + super().__init__() + attention_dim = encoder_output_size + + if input_layer == "embed": + self.embed = torch.nn.Sequential( + torch.nn.Embedding(vocab_size, attention_dim), + pos_enc_class(attention_dim, positional_dropout_rate), + ) + elif input_layer == "linear": + self.embed = torch.nn.Sequential( + torch.nn.Linear(vocab_size, attention_dim), + torch.nn.LayerNorm(attention_dim), + torch.nn.Dropout(dropout_rate), + torch.nn.ReLU(), + pos_enc_class(attention_dim, positional_dropout_rate), + ) + else: + raise ValueError(f"only 'embed' or 'linear' is supported: {input_layer}") + + self.normalize_before = normalize_before + if self.normalize_before: + self.after_norm = LayerNorm(attention_dim) + if use_output_layer: + self.output_layer = torch.nn.Linear(attention_dim, vocab_size) + else: + self.output_layer = None + + # Must set by the inheritance + self.decoders = None + + def forward( + self, + hs_pad: torch.Tensor, + hlens: torch.Tensor, + ys_in_pad: torch.Tensor, + ys_in_lens: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Forward decoder. + + Args: + hs_pad: encoded memory, float32 (batch, maxlen_in, feat) + hlens: (batch) + ys_in_pad: + input token ids, int64 (batch, maxlen_out) + if input_layer == "embed" + input tensor (batch, maxlen_out, #mels) in the other cases + ys_in_lens: (batch) + Returns: + (tuple): tuple containing: + + x: decoded token score before softmax (batch, maxlen_out, token) + if use_output_layer is True, + olens: (batch, ) + """ + tgt = ys_in_pad + # tgt_mask: (B, 1, L) + tgt_mask = (~make_pad_mask(ys_in_lens)[:, None, :]).to(tgt.device) + # m: (1, L, L) + m = subsequent_mask(tgt_mask.size(-1), device=tgt_mask.device).unsqueeze(0) + # tgt_mask: (B, L, L) + tgt_mask = tgt_mask & m + + memory = hs_pad + memory_mask = (~make_pad_mask(hlens, maxlen=memory.size(1)))[:, None, :].to(memory.device) + # Padding for Longformer + if memory_mask.shape[-1] != memory.shape[1]: + padlen = memory.shape[1] - memory_mask.shape[-1] + memory_mask = torch.nn.functional.pad(memory_mask, (0, padlen), "constant", False) + + x = self.embed(tgt) + x, tgt_mask, memory, memory_mask = self.decoders(x, tgt_mask, memory, memory_mask) + if self.normalize_before: + x = self.after_norm(x) + if self.output_layer is not None: + x = self.output_layer(x) + + olens = tgt_mask.sum(1) + return x, olens + + def forward_one_step( + self, + tgt: torch.Tensor, + tgt_mask: torch.Tensor, + memory: torch.Tensor, + cache: List[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, List[torch.Tensor]]: + """Forward one step. + + Args: + tgt: input token ids, int64 (batch, maxlen_out) + tgt_mask: input token mask, (batch, maxlen_out) + dtype=torch.uint8 in PyTorch 1.2- + dtype=torch.bool in PyTorch 1.2+ (include 1.2) + memory: encoded memory, float32 (batch, maxlen_in, feat) + cache: cached output list of (batch, max_time_out-1, size) + Returns: + y, cache: NN output value and cache per `self.decoders`. + y.shape` is (batch, maxlen_out, token) + """ + x = self.embed(tgt) + if cache is None: + cache = [None] * len(self.decoders) + new_cache = [] + for c, decoder in zip(cache, self.decoders): + x, tgt_mask, memory, memory_mask = decoder(x, tgt_mask, memory, None, cache=c) + new_cache.append(x) + + if self.normalize_before: + y = self.after_norm(x[:, -1]) + else: + y = x[:, -1] + if self.output_layer is not None: + y = torch.log_softmax(self.output_layer(y), dim=-1) + + return y, new_cache + + def score(self, ys, state, x): + """Score.""" + ys_mask = subsequent_mask(len(ys), device=x.device).unsqueeze(0) + logp, state = self.forward_one_step(ys.unsqueeze(0), ys_mask, x.unsqueeze(0), cache=state) + return logp.squeeze(0), state + + def batch_score( + self, ys: torch.Tensor, states: List[Any], xs: torch.Tensor + ) -> Tuple[torch.Tensor, List[Any]]: + """Score new token batch. + + Args: + ys (torch.Tensor): torch.int64 prefix tokens (n_batch, ylen). + states (List[Any]): Scorer states for prefix tokens. + xs (torch.Tensor): + The encoder feature that generates ys (n_batch, xlen, n_feat). + + Returns: + tuple[torch.Tensor, List[Any]]: Tuple of + batchfied scores for next token with shape of `(n_batch, n_vocab)` + and next state list for ys. + + """ + # merge states + n_batch = len(ys) + n_layers = len(self.decoders) + if states[0] is None: + batch_state = None + else: + # transpose state of [batch, layer] into [layer, batch] + batch_state = [ + torch.stack([states[b][i] for b in range(n_batch)]) for i in range(n_layers) + ] + + # batch decoding + ys_mask = subsequent_mask(ys.size(-1), device=xs.device).unsqueeze(0) + logp, states = self.forward_one_step(ys, ys_mask, xs, cache=batch_state) + + # transpose state of [layer, batch] into [batch, layer] + state_list = [[states[i][b] for i in range(n_layers)] for b in range(n_batch)] + return logp, state_list + + +@tables.register("decoder_classes", "TransformerDecoder") +class TransformerDecoder(BaseTransformerDecoder): + def __init__( + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + ): + super().__init__( + vocab_size=vocab_size, + encoder_output_size=encoder_output_size, + dropout_rate=dropout_rate, + positional_dropout_rate=positional_dropout_rate, + input_layer=input_layer, + use_output_layer=use_output_layer, + pos_enc_class=pos_enc_class, + normalize_before=normalize_before, + ) + + attention_dim = encoder_output_size + self.decoders = repeat( + num_blocks, + lambda lnum: DecoderLayer( + attention_dim, + MultiHeadedAttention(attention_heads, attention_dim, self_attention_dropout_rate), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), + PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), + dropout_rate, + normalize_before, + concat_after, + ), + ) + + +@tables.register("decoder_classes", "ParaformerDecoderSAN") +class ParaformerDecoderSAN(BaseTransformerDecoder): + """ + Author: Speech Lab of DAMO Academy, Alibaba Group + Paraformer: Fast and Accurate Parallel Transformer for Non-autoregressive End-to-End Speech Recognition + https://arxiv.org/abs/2006.01713 + """ + + def __init__( + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + embeds_id: int = -1, + ): + super().__init__( + vocab_size=vocab_size, + encoder_output_size=encoder_output_size, + dropout_rate=dropout_rate, + positional_dropout_rate=positional_dropout_rate, + input_layer=input_layer, + use_output_layer=use_output_layer, + pos_enc_class=pos_enc_class, + normalize_before=normalize_before, + ) + + attention_dim = encoder_output_size + self.decoders = repeat( + num_blocks, + lambda lnum: DecoderLayer( + attention_dim, + MultiHeadedAttention(attention_heads, attention_dim, self_attention_dropout_rate), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), + PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), + dropout_rate, + normalize_before, + concat_after, + ), + ) + self.embeds_id = embeds_id + self.attention_dim = attention_dim + + def forward( + self, + hs_pad: torch.Tensor, + hlens: torch.Tensor, + ys_in_pad: torch.Tensor, + ys_in_lens: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Forward decoder. + + Args: + hs_pad: encoded memory, float32 (batch, maxlen_in, feat) + hlens: (batch) + ys_in_pad: + input token ids, int64 (batch, maxlen_out) + if input_layer == "embed" + input tensor (batch, maxlen_out, #mels) in the other cases + ys_in_lens: (batch) + Returns: + (tuple): tuple containing: + + x: decoded token score before softmax (batch, maxlen_out, token) + if use_output_layer is True, + olens: (batch, ) + """ + tgt = ys_in_pad + tgt_mask = (~make_pad_mask(ys_in_lens)[:, None, :]).to(tgt.device) + + memory = hs_pad + memory_mask = (~make_pad_mask(hlens, maxlen=memory.size(1)))[:, None, :].to(memory.device) + # Padding for Longformer + if memory_mask.shape[-1] != memory.shape[1]: + padlen = memory.shape[1] - memory_mask.shape[-1] + memory_mask = torch.nn.functional.pad(memory_mask, (0, padlen), "constant", False) + + # x = self.embed(tgt) + x = tgt + embeds_outputs = None + for layer_id, decoder in enumerate(self.decoders): + x, tgt_mask, memory, memory_mask = decoder(x, tgt_mask, memory, memory_mask) + if layer_id == self.embeds_id: + embeds_outputs = x + if self.normalize_before: + x = self.after_norm(x) + if self.output_layer is not None: + x = self.output_layer(x) + + olens = tgt_mask.sum(1) + if embeds_outputs is not None: + return x, olens, embeds_outputs + else: + return x, olens + + +@tables.register("decoder_classes", "LightweightConvolutionTransformerDecoder") +class LightweightConvolutionTransformerDecoder(BaseTransformerDecoder): + def __init__( + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + conv_wshare: int = 4, + conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), + conv_usebias: int = False, + ): + if len(conv_kernel_length) != num_blocks: + raise ValueError( + "conv_kernel_length must have equal number of values to num_blocks: " + f"{len(conv_kernel_length)} != {num_blocks}" + ) + super().__init__( + vocab_size=vocab_size, + encoder_output_size=encoder_output_size, + dropout_rate=dropout_rate, + positional_dropout_rate=positional_dropout_rate, + input_layer=input_layer, + use_output_layer=use_output_layer, + pos_enc_class=pos_enc_class, + normalize_before=normalize_before, + ) + + attention_dim = encoder_output_size + self.decoders = repeat( + num_blocks, + lambda lnum: DecoderLayer( + attention_dim, + LightweightConvolution( + wshare=conv_wshare, + n_feat=attention_dim, + dropout_rate=self_attention_dropout_rate, + kernel_size=conv_kernel_length[lnum], + use_kernel_mask=True, + use_bias=conv_usebias, + ), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), + PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), + dropout_rate, + normalize_before, + concat_after, + ), + ) + + +@tables.register("decoder_classes", "LightweightConvolution2DTransformerDecoder") +class LightweightConvolution2DTransformerDecoder(BaseTransformerDecoder): + def __init__( + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + conv_wshare: int = 4, + conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), + conv_usebias: int = False, + ): + if len(conv_kernel_length) != num_blocks: + raise ValueError( + "conv_kernel_length must have equal number of values to num_blocks: " + f"{len(conv_kernel_length)} != {num_blocks}" + ) + super().__init__( + vocab_size=vocab_size, + encoder_output_size=encoder_output_size, + dropout_rate=dropout_rate, + positional_dropout_rate=positional_dropout_rate, + input_layer=input_layer, + use_output_layer=use_output_layer, + pos_enc_class=pos_enc_class, + normalize_before=normalize_before, + ) + + attention_dim = encoder_output_size + self.decoders = repeat( + num_blocks, + lambda lnum: DecoderLayer( + attention_dim, + LightweightConvolution2D( + wshare=conv_wshare, + n_feat=attention_dim, + dropout_rate=self_attention_dropout_rate, + kernel_size=conv_kernel_length[lnum], + use_kernel_mask=True, + use_bias=conv_usebias, + ), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), + PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), + dropout_rate, + normalize_before, + concat_after, + ), + ) + + +@tables.register("decoder_classes", "DynamicConvolutionTransformerDecoder") +class DynamicConvolutionTransformerDecoder(BaseTransformerDecoder): + def __init__( + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + conv_wshare: int = 4, + conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), + conv_usebias: int = False, + ): + if len(conv_kernel_length) != num_blocks: + raise ValueError( + "conv_kernel_length must have equal number of values to num_blocks: " + f"{len(conv_kernel_length)} != {num_blocks}" + ) + super().__init__( + vocab_size=vocab_size, + encoder_output_size=encoder_output_size, + dropout_rate=dropout_rate, + positional_dropout_rate=positional_dropout_rate, + input_layer=input_layer, + use_output_layer=use_output_layer, + pos_enc_class=pos_enc_class, + normalize_before=normalize_before, + ) + attention_dim = encoder_output_size + + self.decoders = repeat( + num_blocks, + lambda lnum: DecoderLayer( + attention_dim, + DynamicConvolution( + wshare=conv_wshare, + n_feat=attention_dim, + dropout_rate=self_attention_dropout_rate, + kernel_size=conv_kernel_length[lnum], + use_kernel_mask=True, + use_bias=conv_usebias, + ), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), + PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), + dropout_rate, + normalize_before, + concat_after, + ), + ) + + +@tables.register("decoder_classes", "DynamicConvolution2DTransformerDecoder") +class DynamicConvolution2DTransformerDecoder(BaseTransformerDecoder): + def __init__( + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + conv_wshare: int = 4, + conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), + conv_usebias: int = False, + ): + if len(conv_kernel_length) != num_blocks: + raise ValueError( + "conv_kernel_length must have equal number of values to num_blocks: " + f"{len(conv_kernel_length)} != {num_blocks}" + ) + super().__init__( + vocab_size=vocab_size, + encoder_output_size=encoder_output_size, + dropout_rate=dropout_rate, + positional_dropout_rate=positional_dropout_rate, + input_layer=input_layer, + use_output_layer=use_output_layer, + pos_enc_class=pos_enc_class, + normalize_before=normalize_before, + ) + attention_dim = encoder_output_size + + self.decoders = repeat( + num_blocks, + lambda lnum: DecoderLayer( + attention_dim, + DynamicConvolution2D( + wshare=conv_wshare, + n_feat=attention_dim, + dropout_rate=self_attention_dropout_rate, + kernel_size=conv_kernel_length[lnum], + use_kernel_mask=True, + use_bias=conv_usebias, + ), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), + PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), + dropout_rate, + normalize_before, + concat_after, + ), + ) diff --git a/funasr/models/sanm/__init__.py b/funasr/models/sanm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/models/sanm/attention.py b/funasr/models/sanm/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..47d60cb67b4a88374a2f743d4352926e24578e8b --- /dev/null +++ b/funasr/models/sanm/attention.py @@ -0,0 +1,903 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +# Copyright 2019 Shigeki Karita +# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +"""Multi-Head Attention layer definition.""" + +import math + +import numpy +import torch +from torch import nn +from typing import Optional, Tuple + +import torch.nn.functional as F +from funasr.models.transformer.utils.nets_utils import make_pad_mask +import funasr.models.lora.layers as lora + + +def preprocess_for_attn(x, mask, cache, pad_fn, kernel_size): + x = x * mask + x = x.transpose(1, 2) + if cache is None: + x = pad_fn(x) + else: + x = torch.cat((cache, x), dim=2) + cache = x[:, :, -(kernel_size - 1) :] + return x, cache + + +torch_version = tuple([int(i) for i in torch.__version__.split(".")[:2]]) +if torch_version >= (1, 8): + import torch.fx + + torch.fx.wrap("preprocess_for_attn") + + +class MultiHeadedAttention(nn.Module): + """Multi-Head Attention layer. + + Args: + n_head (int): The number of heads. + n_feat (int): The number of features. + dropout_rate (float): Dropout rate. + + """ + + def __init__(self, n_head, n_feat, dropout_rate): + """Construct an MultiHeadedAttention object.""" + super(MultiHeadedAttention, self).__init__() + assert n_feat % n_head == 0 + # We assume d_v always equals d_k + self.d_k = n_feat // n_head + self.h = n_head + self.linear_q = nn.Linear(n_feat, n_feat) + self.linear_k = nn.Linear(n_feat, n_feat) + self.linear_v = nn.Linear(n_feat, n_feat) + self.linear_out = nn.Linear(n_feat, n_feat) + self.attn = None + self.dropout = nn.Dropout(p=dropout_rate) + + def forward_qkv(self, query, key, value): + """Transform query, key and value. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + + Returns: + torch.Tensor: Transformed query tensor (#batch, n_head, time1, d_k). + torch.Tensor: Transformed key tensor (#batch, n_head, time2, d_k). + torch.Tensor: Transformed value tensor (#batch, n_head, time2, d_k). + + """ + n_batch = query.size(0) + q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k) + k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k) + v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k) + q = q.transpose(1, 2) # (batch, head, time1, d_k) + k = k.transpose(1, 2) # (batch, head, time2, d_k) + v = v.transpose(1, 2) # (batch, head, time2, d_k) + + return q, k, v + + def forward_attention(self, value, scores, mask): + """Compute attention context vector. + + Args: + value (torch.Tensor): Transformed value (#batch, n_head, time2, d_k). + scores (torch.Tensor): Attention score (#batch, n_head, time1, time2). + mask (torch.Tensor): Mask (#batch, 1, time2) or (#batch, time1, time2). + + Returns: + torch.Tensor: Transformed value (#batch, time1, d_model) + weighted by the attention score (#batch, time1, time2). + + """ + n_batch = value.size(0) + if mask is not None: + mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) + min_value = -float( + "inf" + ) # float(numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min) + scores = scores.masked_fill(mask, min_value) + attn = torch.softmax(scores, dim=-1).masked_fill( + mask, 0.0 + ) # (batch, head, time1, time2) + else: + attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2) + + p_attn = self.dropout(attn) + x = torch.matmul(p_attn, value) # (batch, head, time1, d_k) + x = ( + x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) + ) # (batch, time1, d_model) + + return self.linear_out(x) # (batch, time1, d_model) + + def forward(self, query, key, value, mask): + """Compute scaled dot product attention. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + mask (torch.Tensor): Mask tensor (#batch, 1, time2) or + (#batch, time1, time2). + + Returns: + torch.Tensor: Output tensor (#batch, time1, d_model). + + """ + q, k, v = self.forward_qkv(query, key, value) + scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) + return self.forward_attention(v, scores, mask) + + +class MultiHeadedAttentionSANM(nn.Module): + """Multi-Head Attention layer. + + Args: + n_head (int): The number of heads. + n_feat (int): The number of features. + dropout_rate (float): Dropout rate. + + """ + + def __init__( + self, + n_head, + in_feat, + n_feat, + dropout_rate, + kernel_size, + sanm_shfit=0, + lora_list=None, + lora_rank=8, + lora_alpha=16, + lora_dropout=0.1, + ): + """Construct an MultiHeadedAttention object.""" + super().__init__() + assert n_feat % n_head == 0 + # We assume d_v always equals d_k + self.d_k = n_feat // n_head + self.h = n_head + # self.linear_q = nn.Linear(n_feat, n_feat) + # self.linear_k = nn.Linear(n_feat, n_feat) + # self.linear_v = nn.Linear(n_feat, n_feat) + if lora_list is not None: + if "o" in lora_list: + self.linear_out = lora.Linear( + n_feat, n_feat, r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout + ) + else: + self.linear_out = nn.Linear(n_feat, n_feat) + lora_qkv_list = ["q" in lora_list, "k" in lora_list, "v" in lora_list] + if lora_qkv_list == [False, False, False]: + self.linear_q_k_v = nn.Linear(in_feat, n_feat * 3) + else: + self.linear_q_k_v = lora.MergedLinear( + in_feat, + n_feat * 3, + r=lora_rank, + lora_alpha=lora_alpha, + lora_dropout=lora_dropout, + enable_lora=lora_qkv_list, + ) + else: + self.linear_out = nn.Linear(n_feat, n_feat) + self.linear_q_k_v = nn.Linear(in_feat, n_feat * 3) + attn = None + self.dropout = nn.Dropout(p=dropout_rate) + + self.fsmn_block = nn.Conv1d( + n_feat, n_feat, kernel_size, stride=1, padding=0, groups=n_feat, bias=False + ) + # padding + left_padding = (kernel_size - 1) // 2 + if sanm_shfit > 0: + left_padding = left_padding + sanm_shfit + right_padding = kernel_size - 1 - left_padding + self.pad_fn = nn.ConstantPad1d((left_padding, right_padding), 0.0) + + def forward_fsmn(self, inputs, mask, mask_shfit_chunk=None): + b, t, d = inputs.size() + if mask is not None: + mask = torch.reshape(mask, (b, -1, 1)) + if mask_shfit_chunk is not None: + mask = mask * mask_shfit_chunk + inputs = inputs * mask + + x = inputs.transpose(1, 2) + x = self.pad_fn(x) + x = self.fsmn_block(x) + x = x.transpose(1, 2) + x += inputs + x = self.dropout(x) + if mask is not None: + x = x * mask + return x + + def forward_qkv(self, x): + """Transform query, key and value. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + + Returns: + torch.Tensor: Transformed query tensor (#batch, n_head, time1, d_k). + torch.Tensor: Transformed key tensor (#batch, n_head, time2, d_k). + torch.Tensor: Transformed value tensor (#batch, n_head, time2, d_k). + + """ + b, t, d = x.size() + q_k_v = self.linear_q_k_v(x) + q, k, v = torch.split(q_k_v, int(self.h * self.d_k), dim=-1) + q_h = torch.reshape(q, (b, t, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time1, d_k) + k_h = torch.reshape(k, (b, t, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time2, d_k) + v_h = torch.reshape(v, (b, t, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time2, d_k) + + return q_h, k_h, v_h, v + + def forward_attention(self, value, scores, mask, mask_att_chunk_encoder=None): + """Compute attention context vector. + + Args: + value (torch.Tensor): Transformed value (#batch, n_head, time2, d_k). + scores (torch.Tensor): Attention score (#batch, n_head, time1, time2). + mask (torch.Tensor): Mask (#batch, 1, time2) or (#batch, time1, time2). + + Returns: + torch.Tensor: Transformed value (#batch, time1, d_model) + weighted by the attention score (#batch, time1, time2). + + """ + n_batch = value.size(0) + if mask is not None: + if mask_att_chunk_encoder is not None: + mask = mask * mask_att_chunk_encoder + + mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) + + min_value = -float( + "inf" + ) # float(numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min) + scores = scores.masked_fill(mask, min_value) + attn = torch.softmax(scores, dim=-1).masked_fill( + mask, 0.0 + ) # (batch, head, time1, time2) + else: + attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2) + + p_attn = self.dropout(attn) + x = torch.matmul(p_attn, value) # (batch, head, time1, d_k) + x = ( + x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) + ) # (batch, time1, d_model) + + return self.linear_out(x) # (batch, time1, d_model) + + def forward(self, x, mask, mask_shfit_chunk=None, mask_att_chunk_encoder=None): + """Compute scaled dot product attention. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + mask (torch.Tensor): Mask tensor (#batch, 1, time2) or + (#batch, time1, time2). + + Returns: + torch.Tensor: Output tensor (#batch, time1, d_model). + + """ + q_h, k_h, v_h, v = self.forward_qkv(x) + fsmn_memory = self.forward_fsmn(v, mask, mask_shfit_chunk) + q_h = q_h * self.d_k ** (-0.5) + scores = torch.matmul(q_h, k_h.transpose(-2, -1)) + att_outs = self.forward_attention(v_h, scores, mask, mask_att_chunk_encoder) + return att_outs + fsmn_memory + + def forward_chunk(self, x, cache=None, chunk_size=None, look_back=0): + """Compute scaled dot product attention. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + mask (torch.Tensor): Mask tensor (#batch, 1, time2) or + (#batch, time1, time2). + + Returns: + torch.Tensor: Output tensor (#batch, time1, d_model). + + """ + q_h, k_h, v_h, v = self.forward_qkv(x) + if chunk_size is not None and look_back > 0 or look_back == -1: + if cache is not None: + k_h_stride = k_h[:, :, : -(chunk_size[2]), :] + v_h_stride = v_h[:, :, : -(chunk_size[2]), :] + k_h = torch.cat((cache["k"], k_h), dim=2) + v_h = torch.cat((cache["v"], v_h), dim=2) + + cache["k"] = torch.cat((cache["k"], k_h_stride), dim=2) + cache["v"] = torch.cat((cache["v"], v_h_stride), dim=2) + if look_back != -1: + cache["k"] = cache["k"][:, :, -(look_back * chunk_size[1]) :, :] + cache["v"] = cache["v"][:, :, -(look_back * chunk_size[1]) :, :] + else: + cache_tmp = { + "k": k_h[:, :, : -(chunk_size[2]), :], + "v": v_h[:, :, : -(chunk_size[2]), :], + } + cache = cache_tmp + fsmn_memory = self.forward_fsmn(v, None) + q_h = q_h * self.d_k ** (-0.5) + scores = torch.matmul(q_h, k_h.transpose(-2, -1)) + att_outs = self.forward_attention(v_h, scores, None) + return att_outs + fsmn_memory, cache + + +class MultiHeadedAttentionSANMExport(nn.Module): + def __init__(self, model): + super().__init__() + self.d_k = model.d_k + self.h = model.h + self.linear_out = model.linear_out + self.linear_q_k_v = model.linear_q_k_v + self.fsmn_block = model.fsmn_block + self.pad_fn = model.pad_fn + + self.attn = None + self.all_head_size = self.h * self.d_k + + def forward(self, x, mask): + mask_3d_btd, mask_4d_bhlt = mask + q_h, k_h, v_h, v = self.forward_qkv(x) + fsmn_memory = self.forward_fsmn(v, mask_3d_btd) + q_h = q_h * self.d_k ** (-0.5) + scores = torch.matmul(q_h, k_h.transpose(-2, -1)) + att_outs = self.forward_attention(v_h, scores, mask_4d_bhlt) + return att_outs + fsmn_memory + + def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor: + new_x_shape = x.size()[:-1] + (self.h, self.d_k) + x = x.view(new_x_shape) + return x.permute(0, 2, 1, 3) + + def forward_qkv(self, x): + q_k_v = self.linear_q_k_v(x) + q, k, v = torch.split(q_k_v, int(self.h * self.d_k), dim=-1) + q_h = self.transpose_for_scores(q) + k_h = self.transpose_for_scores(k) + v_h = self.transpose_for_scores(v) + return q_h, k_h, v_h, v + + def forward_fsmn(self, inputs, mask): + # b, t, d = inputs.size() + # mask = torch.reshape(mask, (b, -1, 1)) + inputs = inputs * mask + x = inputs.transpose(1, 2) + x = self.pad_fn(x) + x = self.fsmn_block(x) + x = x.transpose(1, 2) + x = x + inputs + x = x * mask + return x + + def forward_attention(self, value, scores, mask): + scores = scores + mask + + attn = torch.softmax(scores, dim=-1) + context_layer = torch.matmul(attn, value) # (batch, head, time1, d_k) + + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.view(new_context_layer_shape) + return self.linear_out(context_layer) # (batch, time1, d_model) + + +class MultiHeadedAttentionSANMExport(nn.Module): + def __init__(self, model): + super().__init__() + self.d_k = model.d_k + self.h = model.h + self.linear_out = model.linear_out + self.linear_q_k_v = model.linear_q_k_v + self.fsmn_block = model.fsmn_block + self.pad_fn = model.pad_fn + + self.attn = None + self.all_head_size = self.h * self.d_k + + def forward(self, x, mask): + mask_3d_btd, mask_4d_bhlt = mask + q_h, k_h, v_h, v = self.forward_qkv(x) + fsmn_memory = self.forward_fsmn(v, mask_3d_btd) + q_h = q_h * self.d_k ** (-0.5) + scores = torch.matmul(q_h, k_h.transpose(-2, -1)) + att_outs = self.forward_attention(v_h, scores, mask_4d_bhlt) + return att_outs + fsmn_memory + + def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor: + new_x_shape = x.size()[:-1] + (self.h, self.d_k) + x = x.view(new_x_shape) + return x.permute(0, 2, 1, 3) + + def forward_qkv(self, x): + q_k_v = self.linear_q_k_v(x) + q, k, v = torch.split(q_k_v, int(self.h * self.d_k), dim=-1) + q_h = self.transpose_for_scores(q) + k_h = self.transpose_for_scores(k) + v_h = self.transpose_for_scores(v) + return q_h, k_h, v_h, v + + def forward_fsmn(self, inputs, mask): + # b, t, d = inputs.size() + # mask = torch.reshape(mask, (b, -1, 1)) + inputs = inputs * mask + x = inputs.transpose(1, 2) + x = self.pad_fn(x) + x = self.fsmn_block(x) + x = x.transpose(1, 2) + x = x + inputs + x = x * mask + return x + + def forward_attention(self, value, scores, mask): + scores = scores + mask + + attn = torch.softmax(scores, dim=-1) + context_layer = torch.matmul(attn, value) # (batch, head, time1, d_k) + + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.view(new_context_layer_shape) + return self.linear_out(context_layer) # (batch, time1, d_model) + + +class MultiHeadedAttentionSANMDecoder(nn.Module): + """Multi-Head Attention layer. + + Args: + n_head (int): The number of heads. + n_feat (int): The number of features. + dropout_rate (float): Dropout rate. + + """ + + def __init__(self, n_feat, dropout_rate, kernel_size, sanm_shfit=0): + """Construct an MultiHeadedAttention object.""" + super().__init__() + + self.dropout = nn.Dropout(p=dropout_rate) + + self.fsmn_block = nn.Conv1d( + n_feat, n_feat, kernel_size, stride=1, padding=0, groups=n_feat, bias=False + ) + # padding + # padding + left_padding = (kernel_size - 1) // 2 + if sanm_shfit > 0: + left_padding = left_padding + sanm_shfit + right_padding = kernel_size - 1 - left_padding + self.pad_fn = nn.ConstantPad1d((left_padding, right_padding), 0.0) + self.kernel_size = kernel_size + + def forward(self, inputs, mask, cache=None, mask_shfit_chunk=None): + """ + :param x: (#batch, time1, size). + :param mask: Mask tensor (#batch, 1, time) + :return: + """ + # print("in fsmn, inputs", inputs.size()) + b, t, d = inputs.size() + # logging.info( + # "mask: {}".format(mask.size())) + if mask is not None: + mask = torch.reshape(mask, (b, -1, 1)) + # logging.info("in fsmn, mask: {}, {}".format(mask.size(), mask[0:100:50, :, :])) + if mask_shfit_chunk is not None: + # logging.info("in fsmn, mask_fsmn: {}, {}".format(mask_shfit_chunk.size(), mask_shfit_chunk[0:100:50, :, :])) + mask = mask * mask_shfit_chunk + # logging.info("in fsmn, mask_after_fsmn: {}, {}".format(mask.size(), mask[0:100:50, :, :])) + # print("in fsmn, mask", mask.size()) + # print("in fsmn, inputs", inputs.size()) + inputs = inputs * mask + + x = inputs.transpose(1, 2) + b, d, t = x.size() + if cache is None: + # print("in fsmn, cache is None, x", x.size()) + + x = self.pad_fn(x) + if not self.training: + cache = x + else: + # print("in fsmn, cache is not None, x", x.size()) + # x = torch.cat((x, cache), dim=2)[:, :, :-1] + # if t < self.kernel_size: + # x = self.pad_fn(x) + x = torch.cat((cache[:, :, 1:], x), dim=2) + x = x[:, :, -(self.kernel_size + t - 1) :] + # print("in fsmn, cache is not None, x_cat", x.size()) + cache = x + x = self.fsmn_block(x) + x = x.transpose(1, 2) + # print("in fsmn, fsmn_out", x.size()) + if x.size(1) != inputs.size(1): + inputs = inputs[:, -1, :] + + x = x + inputs + x = self.dropout(x) + if mask is not None: + x = x * mask + return x, cache + + +class MultiHeadedAttentionSANMDecoderExport(nn.Module): + def __init__(self, model): + super().__init__() + self.fsmn_block = model.fsmn_block + self.pad_fn = model.pad_fn + self.kernel_size = model.kernel_size + self.attn = None + + def forward(self, inputs, mask, cache=None): + x, cache = preprocess_for_attn(inputs, mask, cache, self.pad_fn, self.kernel_size) + x = self.fsmn_block(x) + x = x.transpose(1, 2) + + x = x + inputs + x = x * mask + return x, cache + + +class MultiHeadedAttentionCrossAtt(nn.Module): + """Multi-Head Attention layer. + + Args: + n_head (int): The number of heads. + n_feat (int): The number of features. + dropout_rate (float): Dropout rate. + + """ + + def __init__( + self, + n_head, + n_feat, + dropout_rate, + lora_list=None, + lora_rank=8, + lora_alpha=16, + lora_dropout=0.1, + encoder_output_size=None, + ): + """Construct an MultiHeadedAttention object.""" + super().__init__() + assert n_feat % n_head == 0 + # We assume d_v always equals d_k + self.d_k = n_feat // n_head + self.h = n_head + if lora_list is not None: + if "q" in lora_list: + self.linear_q = lora.Linear( + n_feat, n_feat, r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout + ) + else: + self.linear_q = nn.Linear(n_feat, n_feat) + lora_kv_list = ["k" in lora_list, "v" in lora_list] + if lora_kv_list == [False, False]: + self.linear_k_v = nn.Linear( + n_feat if encoder_output_size is None else encoder_output_size, n_feat * 2 + ) + else: + self.linear_k_v = lora.MergedLinear( + n_feat if encoder_output_size is None else encoder_output_size, + n_feat * 2, + r=lora_rank, + lora_alpha=lora_alpha, + lora_dropout=lora_dropout, + enable_lora=lora_kv_list, + ) + if "o" in lora_list: + self.linear_out = lora.Linear( + n_feat, n_feat, r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout + ) + else: + self.linear_out = nn.Linear(n_feat, n_feat) + else: + self.linear_q = nn.Linear(n_feat, n_feat) + self.linear_k_v = nn.Linear( + n_feat if encoder_output_size is None else encoder_output_size, n_feat * 2 + ) + self.linear_out = nn.Linear(n_feat, n_feat) + self.attn = None + self.dropout = nn.Dropout(p=dropout_rate) + + def forward_qkv(self, x, memory): + """Transform query, key and value. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + + Returns: + torch.Tensor: Transformed query tensor (#batch, n_head, time1, d_k). + torch.Tensor: Transformed key tensor (#batch, n_head, time2, d_k). + torch.Tensor: Transformed value tensor (#batch, n_head, time2, d_k). + + """ + + # print("in forward_qkv, x", x.size()) + b = x.size(0) + q = self.linear_q(x) + q_h = torch.reshape(q, (b, -1, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time1, d_k) + + k_v = self.linear_k_v(memory) + k, v = torch.split(k_v, int(self.h * self.d_k), dim=-1) + k_h = torch.reshape(k, (b, -1, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time2, d_k) + v_h = torch.reshape(v, (b, -1, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time2, d_k) + + return q_h, k_h, v_h + + def forward_attention(self, value, scores, mask, ret_attn=False): + """Compute attention context vector. + + Args: + value (torch.Tensor): Transformed value (#batch, n_head, time2, d_k). + scores (torch.Tensor): Attention score (#batch, n_head, time1, time2). + mask (torch.Tensor): Mask (#batch, 1, time2) or (#batch, time1, time2). + + Returns: + torch.Tensor: Transformed value (#batch, time1, d_model) + weighted by the attention score (#batch, time1, time2). + + """ + n_batch = value.size(0) + if mask is not None: + mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) + min_value = -float( + "inf" + ) # float(numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min) + # logging.info( + # "scores: {}, mask_size: {}".format(scores.size(), mask.size())) + scores = scores.masked_fill(mask, min_value) + attn = torch.softmax(scores, dim=-1).masked_fill( + mask, 0.0 + ) # (batch, head, time1, time2) + else: + attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2) + p_attn = self.dropout(attn) + x = torch.matmul(p_attn, value) # (batch, head, time1, d_k) + x = ( + x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) + ) # (batch, time1, d_model) + if ret_attn: + return self.linear_out(x), attn # (batch, time1, d_model) + return self.linear_out(x) # (batch, time1, d_model) + + def forward(self, x, memory, memory_mask, ret_attn=False): + """Compute scaled dot product attention. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + mask (torch.Tensor): Mask tensor (#batch, 1, time2) or + (#batch, time1, time2). + + Returns: + torch.Tensor: Output tensor (#batch, time1, d_model). + + """ + q_h, k_h, v_h = self.forward_qkv(x, memory) + q_h = q_h * self.d_k ** (-0.5) + scores = torch.matmul(q_h, k_h.transpose(-2, -1)) + return self.forward_attention(v_h, scores, memory_mask, ret_attn=ret_attn) + + def forward_chunk(self, x, memory, cache=None, chunk_size=None, look_back=0): + """Compute scaled dot product attention. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + mask (torch.Tensor): Mask tensor (#batch, 1, time2) or + (#batch, time1, time2). + + Returns: + torch.Tensor: Output tensor (#batch, time1, d_model). + + """ + q_h, k_h, v_h = self.forward_qkv(x, memory) + if chunk_size is not None and look_back > 0: + if cache is not None: + k_h = torch.cat((cache["k"], k_h), dim=2) + v_h = torch.cat((cache["v"], v_h), dim=2) + cache["k"] = k_h[:, :, -(look_back * chunk_size[1]) :, :] + cache["v"] = v_h[:, :, -(look_back * chunk_size[1]) :, :] + else: + cache_tmp = { + "k": k_h[:, :, -(look_back * chunk_size[1]) :, :], + "v": v_h[:, :, -(look_back * chunk_size[1]) :, :], + } + cache = cache_tmp + q_h = q_h * self.d_k ** (-0.5) + scores = torch.matmul(q_h, k_h.transpose(-2, -1)) + return self.forward_attention(v_h, scores, None), cache + + +class MultiHeadedAttentionCrossAttExport(nn.Module): + def __init__(self, model): + super().__init__() + self.d_k = model.d_k + self.h = model.h + self.linear_q = model.linear_q + self.linear_k_v = model.linear_k_v + self.linear_out = model.linear_out + self.attn = None + self.all_head_size = self.h * self.d_k + + def forward(self, x, memory, memory_mask, ret_attn=False): + q, k, v = self.forward_qkv(x, memory) + scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) + return self.forward_attention(v, scores, memory_mask, ret_attn) + + def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor: + new_x_shape = x.size()[:-1] + (self.h, self.d_k) + x = x.view(new_x_shape) + return x.permute(0, 2, 1, 3) + + def forward_qkv(self, x, memory): + q = self.linear_q(x) + + k_v = self.linear_k_v(memory) + k, v = torch.split(k_v, int(self.h * self.d_k), dim=-1) + q = self.transpose_for_scores(q) + k = self.transpose_for_scores(k) + v = self.transpose_for_scores(v) + return q, k, v + + def forward_attention(self, value, scores, mask, ret_attn): + scores = scores + mask.to(scores.device) + + attn = torch.softmax(scores, dim=-1) + context_layer = torch.matmul(attn, value) # (batch, head, time1, d_k) + + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.view(new_context_layer_shape) + if ret_attn: + return self.linear_out(context_layer), attn + return self.linear_out(context_layer) # (batch, time1, d_model) + + +class MultiHeadSelfAttention(nn.Module): + """Multi-Head Attention layer. + + Args: + n_head (int): The number of heads. + n_feat (int): The number of features. + dropout_rate (float): Dropout rate. + + """ + + def __init__(self, n_head, in_feat, n_feat, dropout_rate): + """Construct an MultiHeadedAttention object.""" + super(MultiHeadSelfAttention, self).__init__() + assert n_feat % n_head == 0 + # We assume d_v always equals d_k + self.d_k = n_feat // n_head + self.h = n_head + self.linear_out = nn.Linear(n_feat, n_feat) + self.linear_q_k_v = nn.Linear(in_feat, n_feat * 3) + self.attn = None + self.dropout = nn.Dropout(p=dropout_rate) + + def forward_qkv(self, x): + """Transform query, key and value. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + + Returns: + torch.Tensor: Transformed query tensor (#batch, n_head, time1, d_k). + torch.Tensor: Transformed key tensor (#batch, n_head, time2, d_k). + torch.Tensor: Transformed value tensor (#batch, n_head, time2, d_k). + + """ + b, t, d = x.size() + q_k_v = self.linear_q_k_v(x) + q, k, v = torch.split(q_k_v, int(self.h * self.d_k), dim=-1) + q_h = torch.reshape(q, (b, t, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time1, d_k) + k_h = torch.reshape(k, (b, t, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time2, d_k) + v_h = torch.reshape(v, (b, t, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time2, d_k) + + return q_h, k_h, v_h, v + + def forward_attention(self, value, scores, mask, mask_att_chunk_encoder=None): + """Compute attention context vector. + + Args: + value (torch.Tensor): Transformed value (#batch, n_head, time2, d_k). + scores (torch.Tensor): Attention score (#batch, n_head, time1, time2). + mask (torch.Tensor): Mask (#batch, 1, time2) or (#batch, time1, time2). + + Returns: + torch.Tensor: Transformed value (#batch, time1, d_model) + weighted by the attention score (#batch, time1, time2). + + """ + n_batch = value.size(0) + if mask is not None: + if mask_att_chunk_encoder is not None: + mask = mask * mask_att_chunk_encoder + + mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) + + min_value = -float( + "inf" + ) # float(numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min) + scores = scores.masked_fill(mask, min_value) + attn = torch.softmax(scores, dim=-1).masked_fill( + mask, 0.0 + ) # (batch, head, time1, time2) + else: + attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2) + + p_attn = self.dropout(attn) + x = torch.matmul(p_attn, value) # (batch, head, time1, d_k) + x = ( + x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) + ) # (batch, time1, d_model) + + return self.linear_out(x) # (batch, time1, d_model) + + def forward(self, x, mask, mask_att_chunk_encoder=None): + """Compute scaled dot product attention. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + mask (torch.Tensor): Mask tensor (#batch, 1, time2) or + (#batch, time1, time2). + + Returns: + torch.Tensor: Output tensor (#batch, time1, d_model). + + """ + q_h, k_h, v_h, v = self.forward_qkv(x) + q_h = q_h * self.d_k ** (-0.5) + scores = torch.matmul(q_h, k_h.transpose(-2, -1)) + att_outs = self.forward_attention(v_h, scores, mask, mask_att_chunk_encoder) + return att_outs diff --git a/funasr/models/sanm/decoder.py b/funasr/models/sanm/decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..1a4fb26e2b1603bbdb3b21865a2fe39ece3316f0 --- /dev/null +++ b/funasr/models/sanm/decoder.py @@ -0,0 +1,498 @@ +#!/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) + +from typing import List +from typing import Tuple +import logging +import torch +import torch.nn as nn +import numpy as np + +from funasr.models.scama import utils as myutils +from funasr.models.transformer.decoder import BaseTransformerDecoder + +from funasr.models.sanm.attention import ( + MultiHeadedAttentionSANMDecoder, + MultiHeadedAttentionCrossAtt, +) +from funasr.models.transformer.embedding import PositionalEncoding +from funasr.models.transformer.layer_norm import LayerNorm +from funasr.models.sanm.positionwise_feed_forward import PositionwiseFeedForwardDecoderSANM +from funasr.models.transformer.utils.repeat import repeat + +from funasr.register import tables + + +class DecoderLayerSANM(nn.Module): + """Single decoder layer module. + + Args: + size (int): Input dimension. + self_attn (torch.nn.Module): Self-attention module instance. + `MultiHeadedAttention` instance can be used as the argument. + src_attn (torch.nn.Module): Self-attention module instance. + `MultiHeadedAttention` instance can be used as the argument. + feed_forward (torch.nn.Module): Feed-forward module instance. + `PositionwiseFeedForward`, `MultiLayeredConv1d`, or `Conv1dLinear` instance + can be used as the argument. + dropout_rate (float): Dropout rate. + normalize_before (bool): Whether to use layer_norm before the first block. + concat_after (bool): Whether to concat attention layer's input and output. + if True, additional linear will be applied. + i.e. x -> x + linear(concat(x, att(x))) + if False, no additional linear will be applied. i.e. x -> x + att(x) + + + """ + + def __init__( + self, + size, + self_attn, + src_attn, + feed_forward, + dropout_rate, + normalize_before=True, + concat_after=False, + ): + """Construct an DecoderLayer object.""" + super(DecoderLayerSANM, self).__init__() + self.size = size + self.self_attn = self_attn + self.src_attn = src_attn + self.feed_forward = feed_forward + self.norm1 = LayerNorm(size) + if self_attn is not None: + self.norm2 = LayerNorm(size) + if src_attn is not None: + self.norm3 = LayerNorm(size) + self.dropout = nn.Dropout(dropout_rate) + self.normalize_before = normalize_before + self.concat_after = concat_after + if self.concat_after: + self.concat_linear1 = nn.Linear(size + size, size) + self.concat_linear2 = nn.Linear(size + size, size) + + def forward(self, tgt, tgt_mask, memory, memory_mask=None, cache=None): + """Compute decoded features. + + Args: + tgt (torch.Tensor): Input tensor (#batch, maxlen_out, size). + tgt_mask (torch.Tensor): Mask for input tensor (#batch, maxlen_out). + memory (torch.Tensor): Encoded memory, float32 (#batch, maxlen_in, size). + memory_mask (torch.Tensor): Encoded memory mask (#batch, maxlen_in). + cache (List[torch.Tensor]): List of cached tensors. + Each tensor shape should be (#batch, maxlen_out - 1, size). + + Returns: + torch.Tensor: Output tensor(#batch, maxlen_out, size). + torch.Tensor: Mask for output tensor (#batch, maxlen_out). + torch.Tensor: Encoded memory (#batch, maxlen_in, size). + torch.Tensor: Encoded memory mask (#batch, maxlen_in). + + """ + # tgt = self.dropout(tgt) + residual = tgt + if self.normalize_before: + tgt = self.norm1(tgt) + tgt = self.feed_forward(tgt) + + x = tgt + if self.self_attn: + if self.normalize_before: + tgt = self.norm2(tgt) + x, _ = self.self_attn(tgt, tgt_mask) + x = residual + self.dropout(x) + + if self.src_attn is not None: + residual = x + if self.normalize_before: + x = self.norm3(x) + + x = residual + self.dropout(self.src_attn(x, memory, memory_mask)) + + return x, tgt_mask, memory, memory_mask, cache + + def forward_one_step(self, tgt, tgt_mask, memory, memory_mask=None, cache=None): + """Compute decoded features. + + Args: + tgt (torch.Tensor): Input tensor (#batch, maxlen_out, size). + tgt_mask (torch.Tensor): Mask for input tensor (#batch, maxlen_out). + memory (torch.Tensor): Encoded memory, float32 (#batch, maxlen_in, size). + memory_mask (torch.Tensor): Encoded memory mask (#batch, maxlen_in). + cache (List[torch.Tensor]): List of cached tensors. + Each tensor shape should be (#batch, maxlen_out - 1, size). + + Returns: + torch.Tensor: Output tensor(#batch, maxlen_out, size). + torch.Tensor: Mask for output tensor (#batch, maxlen_out). + torch.Tensor: Encoded memory (#batch, maxlen_in, size). + torch.Tensor: Encoded memory mask (#batch, maxlen_in). + + """ + # tgt = self.dropout(tgt) + residual = tgt + if self.normalize_before: + tgt = self.norm1(tgt) + tgt = self.feed_forward(tgt) + + x = tgt + if self.self_attn: + if self.normalize_before: + tgt = self.norm2(tgt) + if self.training: + cache = None + x, cache = self.self_attn(tgt, tgt_mask, cache=cache) + x = residual + self.dropout(x) + + if self.src_attn is not None: + residual = x + if self.normalize_before: + x = self.norm3(x) + + x = residual + self.dropout(self.src_attn(x, memory, memory_mask)) + + return x, tgt_mask, memory, memory_mask, cache + + def forward_chunk( + self, tgt, memory, fsmn_cache=None, opt_cache=None, chunk_size=None, look_back=0 + ): + """Compute decoded features. + + Args: + tgt (torch.Tensor): Input tensor (#batch, maxlen_out, size). + tgt_mask (torch.Tensor): Mask for input tensor (#batch, maxlen_out). + memory (torch.Tensor): Encoded memory, float32 (#batch, maxlen_in, size). + memory_mask (torch.Tensor): Encoded memory mask (#batch, maxlen_in). + cache (List[torch.Tensor]): List of cached tensors. + Each tensor shape should be (#batch, maxlen_out - 1, size). + + Returns: + torch.Tensor: Output tensor(#batch, maxlen_out, size). + torch.Tensor: Mask for output tensor (#batch, maxlen_out). + torch.Tensor: Encoded memory (#batch, maxlen_in, size). + torch.Tensor: Encoded memory mask (#batch, maxlen_in). + + """ + residual = tgt + if self.normalize_before: + tgt = self.norm1(tgt) + tgt = self.feed_forward(tgt) + + x = tgt + if self.self_attn: + if self.normalize_before: + tgt = self.norm2(tgt) + x, fsmn_cache = self.self_attn(tgt, None, fsmn_cache) + x = residual + self.dropout(x) + + if self.src_attn is not None: + residual = x + if self.normalize_before: + x = self.norm3(x) + + x, opt_cache = self.src_attn.forward_chunk(x, memory, opt_cache, chunk_size, look_back) + x = residual + x + + return x, memory, fsmn_cache, opt_cache + + +@tables.register("decoder_classes", "FsmnDecoder") +class FsmnDecoder(BaseTransformerDecoder): + """ + Author: Zhifu Gao, Shiliang Zhang, Ming Lei, Ian McLoughlin + San-m: Memory equipped self-attention for end-to-end speech recognition + https://arxiv.org/abs/2006.01713 + """ + + def __init__( + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + att_layer_num: int = 6, + kernel_size: int = 21, + sanm_shfit: int = None, + concat_embeds: bool = False, + attention_dim: int = None, + tf2torch_tensor_name_prefix_torch: str = "decoder", + tf2torch_tensor_name_prefix_tf: str = "seq2seq/decoder", + embed_tensor_name_prefix_tf: str = None, + ): + super().__init__( + vocab_size=vocab_size, + encoder_output_size=encoder_output_size, + dropout_rate=dropout_rate, + positional_dropout_rate=positional_dropout_rate, + input_layer=input_layer, + use_output_layer=use_output_layer, + pos_enc_class=pos_enc_class, + normalize_before=normalize_before, + ) + if attention_dim is None: + attention_dim = encoder_output_size + + if input_layer == "embed": + self.embed = torch.nn.Sequential( + torch.nn.Embedding(vocab_size, attention_dim), + ) + elif input_layer == "linear": + self.embed = torch.nn.Sequential( + torch.nn.Linear(vocab_size, attention_dim), + torch.nn.LayerNorm(attention_dim), + torch.nn.Dropout(dropout_rate), + torch.nn.ReLU(), + pos_enc_class(attention_dim, positional_dropout_rate), + ) + else: + raise ValueError(f"only 'embed' or 'linear' is supported: {input_layer}") + + self.normalize_before = normalize_before + if self.normalize_before: + self.after_norm = LayerNorm(attention_dim) + if use_output_layer: + self.output_layer = torch.nn.Linear(attention_dim, vocab_size) + else: + self.output_layer = None + + self.att_layer_num = att_layer_num + self.num_blocks = num_blocks + if sanm_shfit is None: + sanm_shfit = (kernel_size - 1) // 2 + self.decoders = repeat( + att_layer_num, + lambda lnum: DecoderLayerSANM( + attention_dim, + MultiHeadedAttentionSANMDecoder( + attention_dim, self_attention_dropout_rate, kernel_size, sanm_shfit=sanm_shfit + ), + MultiHeadedAttentionCrossAtt( + attention_heads, + attention_dim, + src_attention_dropout_rate, + encoder_output_size=encoder_output_size, + ), + PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate), + dropout_rate, + normalize_before, + concat_after, + ), + ) + if num_blocks - att_layer_num <= 0: + self.decoders2 = None + else: + self.decoders2 = repeat( + num_blocks - att_layer_num, + lambda lnum: DecoderLayerSANM( + attention_dim, + MultiHeadedAttentionSANMDecoder( + attention_dim, + self_attention_dropout_rate, + kernel_size, + sanm_shfit=sanm_shfit, + ), + None, + PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate), + dropout_rate, + normalize_before, + concat_after, + ), + ) + + self.decoders3 = repeat( + 1, + lambda lnum: DecoderLayerSANM( + attention_dim, + None, + None, + PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate), + dropout_rate, + normalize_before, + concat_after, + ), + ) + if concat_embeds: + self.embed_concat_ffn = repeat( + 1, + lambda lnum: DecoderLayerSANM( + attention_dim + encoder_output_size, + None, + None, + PositionwiseFeedForwardDecoderSANM( + attention_dim + encoder_output_size, + linear_units, + dropout_rate, + adim=attention_dim, + ), + dropout_rate, + normalize_before, + concat_after, + ), + ) + else: + self.embed_concat_ffn = None + self.concat_embeds = concat_embeds + self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch + self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf + self.embed_tensor_name_prefix_tf = embed_tensor_name_prefix_tf + + def forward( + self, + hs_pad: torch.Tensor, + hlens: torch.Tensor, + ys_in_pad: torch.Tensor, + ys_in_lens: torch.Tensor, + chunk_mask: torch.Tensor = None, + pre_acoustic_embeds: torch.Tensor = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Forward decoder. + + Args: + hs_pad: encoded memory, float32 (batch, maxlen_in, feat) + hlens: (batch) + ys_in_pad: + input token ids, int64 (batch, maxlen_out) + if input_layer == "embed" + input tensor (batch, maxlen_out, #mels) in the other cases + ys_in_lens: (batch) + Returns: + (tuple): tuple containing: + + x: decoded token score before softmax (batch, maxlen_out, token) + if use_output_layer is True, + olens: (batch, ) + """ + tgt = ys_in_pad + tgt_mask = myutils.sequence_mask(ys_in_lens, device=tgt.device)[:, :, None] + + memory = hs_pad + memory_mask = myutils.sequence_mask(hlens, device=memory.device)[:, None, :] + if chunk_mask is not None: + memory_mask = memory_mask * chunk_mask + if tgt_mask.size(1) != memory_mask.size(1): + memory_mask = torch.cat((memory_mask, memory_mask[:, -2:-1, :]), dim=1) + + x = self.embed(tgt) + + if pre_acoustic_embeds is not None and self.concat_embeds: + x = torch.cat((x, pre_acoustic_embeds), dim=-1) + x, _, _, _, _ = self.embed_concat_ffn(x, None, None, None, None) + + x, tgt_mask, memory, memory_mask, _ = self.decoders(x, tgt_mask, memory, memory_mask) + if self.decoders2 is not None: + x, tgt_mask, memory, memory_mask, _ = self.decoders2(x, tgt_mask, memory, memory_mask) + x, tgt_mask, memory, memory_mask, _ = self.decoders3(x, tgt_mask, memory, memory_mask) + if self.normalize_before: + x = self.after_norm(x) + if self.output_layer is not None: + x = self.output_layer(x) + + olens = tgt_mask.sum(1) + return x, olens + + def score( + self, + ys, + state, + x, + x_mask=None, + pre_acoustic_embeds: torch.Tensor = None, + ): + """Score.""" + ys_mask = myutils.sequence_mask( + torch.tensor([len(ys)], dtype=torch.int32), device=x.device + )[:, :, None] + logp, state = self.forward_one_step( + ys.unsqueeze(0), + ys_mask, + x.unsqueeze(0), + memory_mask=x_mask, + pre_acoustic_embeds=pre_acoustic_embeds, + cache=state, + ) + return logp.squeeze(0), state + + def forward_one_step( + self, + tgt: torch.Tensor, + tgt_mask: torch.Tensor, + memory: torch.Tensor, + memory_mask: torch.Tensor = None, + pre_acoustic_embeds: torch.Tensor = None, + cache: List[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, List[torch.Tensor]]: + """Forward one step. + + Args: + tgt: input token ids, int64 (batch, maxlen_out) + tgt_mask: input token mask, (batch, maxlen_out) + dtype=torch.uint8 in PyTorch 1.2- + dtype=torch.bool in PyTorch 1.2+ (include 1.2) + memory: encoded memory, float32 (batch, maxlen_in, feat) + cache: cached output list of (batch, max_time_out-1, size) + Returns: + y, cache: NN output value and cache per `self.decoders`. + y.shape` is (batch, maxlen_out, token) + """ + + x = tgt[:, -1:] + tgt_mask = None + x = self.embed(x) + + if pre_acoustic_embeds is not None and self.concat_embeds: + x = torch.cat((x, pre_acoustic_embeds), dim=-1) + x, _, _, _, _ = self.embed_concat_ffn(x, None, None, None, None) + + if cache is None: + cache_layer_num = len(self.decoders) + if self.decoders2 is not None: + cache_layer_num += len(self.decoders2) + cache = [None] * cache_layer_num + new_cache = [] + # for c, decoder in zip(cache, self.decoders): + for i in range(self.att_layer_num): + decoder = self.decoders[i] + c = cache[i] + x, tgt_mask, memory, memory_mask, c_ret = decoder.forward_one_step( + x, tgt_mask, memory, memory_mask, cache=c + ) + new_cache.append(c_ret) + + if self.num_blocks - self.att_layer_num >= 1: + for i in range(self.num_blocks - self.att_layer_num): + j = i + self.att_layer_num + decoder = self.decoders2[i] + c = cache[j] + x, tgt_mask, memory, memory_mask, c_ret = decoder.forward_one_step( + x, tgt_mask, memory, memory_mask, cache=c + ) + new_cache.append(c_ret) + + for decoder in self.decoders3: + x, tgt_mask, memory, memory_mask, _ = decoder.forward_one_step( + x, tgt_mask, memory, None, cache=None + ) + + if self.normalize_before: + y = self.after_norm(x[:, -1]) + else: + y = x[:, -1] + if self.output_layer is not None: + y = self.output_layer(y) + y = torch.log_softmax(y, dim=-1) + + return y, new_cache diff --git a/funasr/models/sanm/encoder.py b/funasr/models/sanm/encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..0d39ca7424c847995991615b2f01ee6800b25995 --- /dev/null +++ b/funasr/models/sanm/encoder.py @@ -0,0 +1,612 @@ +#!/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) + +from typing import List +from typing import Optional +from typing import Sequence +from typing import Tuple +from typing import Union +import logging +import torch +import torch.nn as nn +import torch.nn.functional as F + +import numpy as np +from funasr.train_utils.device_funcs import to_device +from funasr.models.transformer.utils.nets_utils import make_pad_mask +from funasr.models.sanm.attention import MultiHeadedAttention, MultiHeadedAttentionSANM +from funasr.models.transformer.embedding import ( + SinusoidalPositionEncoder, + StreamSinusoidalPositionEncoder, +) +from funasr.models.transformer.layer_norm import LayerNorm +from funasr.models.transformer.utils.multi_layer_conv import Conv1dLinear +from funasr.models.transformer.utils.multi_layer_conv import MultiLayeredConv1d +from funasr.models.transformer.positionwise_feed_forward import ( + PositionwiseFeedForward, # noqa: H301 +) +from funasr.models.transformer.utils.repeat import repeat +from funasr.models.transformer.utils.subsampling import Conv2dSubsampling +from funasr.models.transformer.utils.subsampling import Conv2dSubsampling2 +from funasr.models.transformer.utils.subsampling import Conv2dSubsampling6 +from funasr.models.transformer.utils.subsampling import Conv2dSubsampling8 +from funasr.models.transformer.utils.subsampling import TooShortUttError +from funasr.models.transformer.utils.subsampling import check_short_utt + + +from funasr.models.ctc.ctc import CTC + +from funasr.register import tables + + +class EncoderLayerSANM(nn.Module): + def __init__( + self, + in_size, + size, + self_attn, + feed_forward, + dropout_rate, + normalize_before=True, + concat_after=False, + stochastic_depth_rate=0.0, + ): + """Construct an EncoderLayer object.""" + super(EncoderLayerSANM, self).__init__() + self.self_attn = self_attn + self.feed_forward = feed_forward + self.norm1 = LayerNorm(in_size) + self.norm2 = LayerNorm(size) + self.dropout = nn.Dropout(dropout_rate) + self.in_size = in_size + self.size = size + self.normalize_before = normalize_before + self.concat_after = concat_after + if self.concat_after: + self.concat_linear = nn.Linear(size + size, size) + self.stochastic_depth_rate = stochastic_depth_rate + self.dropout_rate = dropout_rate + + def forward(self, x, mask, cache=None, mask_shfit_chunk=None, mask_att_chunk_encoder=None): + """Compute encoded features. + + Args: + x_input (torch.Tensor): Input tensor (#batch, time, size). + mask (torch.Tensor): Mask tensor for the input (#batch, time). + cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size). + + Returns: + torch.Tensor: Output tensor (#batch, time, size). + torch.Tensor: Mask tensor (#batch, time). + + """ + skip_layer = False + # with stochastic depth, residual connection `x + f(x)` becomes + # `x <- x + 1 / (1 - p) * f(x)` at training time. + stoch_layer_coeff = 1.0 + if self.training and self.stochastic_depth_rate > 0: + skip_layer = torch.rand(1).item() < self.stochastic_depth_rate + stoch_layer_coeff = 1.0 / (1 - self.stochastic_depth_rate) + + if skip_layer: + if cache is not None: + x = torch.cat([cache, x], dim=1) + return x, mask + + residual = x + if self.normalize_before: + x = self.norm1(x) + + if self.concat_after: + x_concat = torch.cat( + ( + x, + self.self_attn( + x, + mask, + mask_shfit_chunk=mask_shfit_chunk, + mask_att_chunk_encoder=mask_att_chunk_encoder, + ), + ), + dim=-1, + ) + if self.in_size == self.size: + x = residual + stoch_layer_coeff * self.concat_linear(x_concat) + else: + x = stoch_layer_coeff * self.concat_linear(x_concat) + else: + if self.in_size == self.size: + x = residual + stoch_layer_coeff * self.dropout( + self.self_attn( + x, + mask, + mask_shfit_chunk=mask_shfit_chunk, + mask_att_chunk_encoder=mask_att_chunk_encoder, + ) + ) + else: + x = stoch_layer_coeff * self.dropout( + self.self_attn( + x, + mask, + mask_shfit_chunk=mask_shfit_chunk, + mask_att_chunk_encoder=mask_att_chunk_encoder, + ) + ) + if not self.normalize_before: + x = self.norm1(x) + + residual = x + if self.normalize_before: + x = self.norm2(x) + x = residual + stoch_layer_coeff * self.dropout(self.feed_forward(x)) + if not self.normalize_before: + x = self.norm2(x) + + return x, mask, cache, mask_shfit_chunk, mask_att_chunk_encoder + + def forward_chunk(self, x, cache=None, chunk_size=None, look_back=0): + """Compute encoded features. + + Args: + x_input (torch.Tensor): Input tensor (#batch, time, size). + mask (torch.Tensor): Mask tensor for the input (#batch, time). + cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size). + + Returns: + torch.Tensor: Output tensor (#batch, time, size). + torch.Tensor: Mask tensor (#batch, time). + + """ + + residual = x + if self.normalize_before: + x = self.norm1(x) + + if self.in_size == self.size: + attn, cache = self.self_attn.forward_chunk(x, cache, chunk_size, look_back) + x = residual + attn + else: + x, cache = self.self_attn.forward_chunk(x, cache, chunk_size, look_back) + + if not self.normalize_before: + x = self.norm1(x) + + residual = x + if self.normalize_before: + x = self.norm2(x) + x = residual + self.feed_forward(x) + if not self.normalize_before: + x = self.norm2(x) + + return x, cache + + +@tables.register("encoder_classes", "SANMEncoder") +class SANMEncoder(nn.Module): + """ + Author: Zhifu Gao, Shiliang Zhang, Ming Lei, Ian McLoughlin + San-m: Memory equipped self-attention for end-to-end speech recognition + https://arxiv.org/abs/2006.01713 + """ + + def __init__( + self, + input_size: int, + output_size: int = 256, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + attention_dropout_rate: float = 0.0, + input_layer: Optional[str] = "conv2d", + pos_enc_class=SinusoidalPositionEncoder, + normalize_before: bool = True, + concat_after: bool = False, + positionwise_layer_type: str = "linear", + positionwise_conv_kernel_size: int = 1, + padding_idx: int = -1, + interctc_layer_idx: List[int] = [], + interctc_use_conditioning: bool = False, + kernel_size: int = 11, + sanm_shfit: int = 0, + lora_list: List[str] = None, + lora_rank: int = 8, + lora_alpha: int = 16, + lora_dropout: float = 0.1, + selfattention_layer_type: str = "sanm", + tf2torch_tensor_name_prefix_torch: str = "encoder", + tf2torch_tensor_name_prefix_tf: str = "seq2seq/encoder", + ): + super().__init__() + self._output_size = output_size + + if input_layer == "linear": + self.embed = torch.nn.Sequential( + torch.nn.Linear(input_size, output_size), + torch.nn.LayerNorm(output_size), + torch.nn.Dropout(dropout_rate), + torch.nn.ReLU(), + pos_enc_class(output_size, positional_dropout_rate), + ) + elif input_layer == "conv2d": + self.embed = Conv2dSubsampling(input_size, output_size, dropout_rate) + elif input_layer == "conv2d2": + self.embed = Conv2dSubsampling2(input_size, output_size, dropout_rate) + elif input_layer == "conv2d6": + self.embed = Conv2dSubsampling6(input_size, output_size, dropout_rate) + elif input_layer == "conv2d8": + self.embed = Conv2dSubsampling8(input_size, output_size, dropout_rate) + elif input_layer == "embed": + self.embed = torch.nn.Sequential( + torch.nn.Embedding(input_size, output_size, padding_idx=padding_idx), + SinusoidalPositionEncoder(), + ) + elif input_layer is None: + if input_size == output_size: + self.embed = None + else: + self.embed = torch.nn.Linear(input_size, output_size) + elif input_layer == "pe": + self.embed = SinusoidalPositionEncoder() + elif input_layer == "pe_online": + self.embed = StreamSinusoidalPositionEncoder() + else: + raise ValueError("unknown input_layer: " + input_layer) + self.normalize_before = normalize_before + if positionwise_layer_type == "linear": + positionwise_layer = PositionwiseFeedForward + positionwise_layer_args = ( + output_size, + linear_units, + dropout_rate, + ) + elif positionwise_layer_type == "conv1d": + positionwise_layer = MultiLayeredConv1d + positionwise_layer_args = ( + output_size, + linear_units, + positionwise_conv_kernel_size, + dropout_rate, + ) + elif positionwise_layer_type == "conv1d-linear": + positionwise_layer = Conv1dLinear + positionwise_layer_args = ( + output_size, + linear_units, + positionwise_conv_kernel_size, + dropout_rate, + ) + else: + raise NotImplementedError("Support only linear or conv1d.") + + if selfattention_layer_type == "selfattn": + encoder_selfattn_layer = MultiHeadedAttention + encoder_selfattn_layer_args = ( + attention_heads, + output_size, + attention_dropout_rate, + ) + + elif selfattention_layer_type == "sanm": + encoder_selfattn_layer = MultiHeadedAttentionSANM + encoder_selfattn_layer_args0 = ( + attention_heads, + input_size, + output_size, + attention_dropout_rate, + kernel_size, + sanm_shfit, + lora_list, + lora_rank, + lora_alpha, + lora_dropout, + ) + + encoder_selfattn_layer_args = ( + attention_heads, + output_size, + output_size, + attention_dropout_rate, + kernel_size, + sanm_shfit, + lora_list, + lora_rank, + lora_alpha, + lora_dropout, + ) + self.encoders0 = repeat( + 1, + lambda lnum: EncoderLayerSANM( + input_size, + output_size, + encoder_selfattn_layer(*encoder_selfattn_layer_args0), + positionwise_layer(*positionwise_layer_args), + dropout_rate, + normalize_before, + concat_after, + ), + ) + + self.encoders = repeat( + num_blocks - 1, + lambda lnum: EncoderLayerSANM( + output_size, + output_size, + encoder_selfattn_layer(*encoder_selfattn_layer_args), + positionwise_layer(*positionwise_layer_args), + dropout_rate, + normalize_before, + concat_after, + ), + ) + if self.normalize_before: + self.after_norm = LayerNorm(output_size) + + self.interctc_layer_idx = interctc_layer_idx + if len(interctc_layer_idx) > 0: + assert 0 < min(interctc_layer_idx) and max(interctc_layer_idx) < num_blocks + self.interctc_use_conditioning = interctc_use_conditioning + self.conditioning_layer = None + self.dropout = nn.Dropout(dropout_rate) + self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch + self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf + + def output_size(self) -> int: + return self._output_size + + def forward( + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + prev_states: torch.Tensor = None, + ctc: CTC = None, + ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Embed positions in tensor. + + Args: + xs_pad: input tensor (B, L, D) + ilens: input length (B) + prev_states: Not to be used now. + Returns: + position embedded tensor and mask + """ + masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device) + xs_pad = xs_pad * self.output_size() ** 0.5 + if self.embed is None: + xs_pad = xs_pad + elif ( + isinstance(self.embed, Conv2dSubsampling) + or isinstance(self.embed, Conv2dSubsampling2) + or isinstance(self.embed, Conv2dSubsampling6) + or isinstance(self.embed, Conv2dSubsampling8) + ): + short_status, limit_size = check_short_utt(self.embed, xs_pad.size(1)) + if short_status: + raise TooShortUttError( + f"has {xs_pad.size(1)} frames and is too short for subsampling " + + f"(it needs more than {limit_size} frames), return empty results", + xs_pad.size(1), + limit_size, + ) + xs_pad, masks = self.embed(xs_pad, masks) + else: + xs_pad = self.embed(xs_pad) + + # xs_pad = self.dropout(xs_pad) + encoder_outs = self.encoders0(xs_pad, masks) + xs_pad, masks = encoder_outs[0], encoder_outs[1] + intermediate_outs = [] + if len(self.interctc_layer_idx) == 0: + encoder_outs = self.encoders(xs_pad, masks) + xs_pad, masks = encoder_outs[0], encoder_outs[1] + else: + for layer_idx, encoder_layer in enumerate(self.encoders): + encoder_outs = encoder_layer(xs_pad, masks) + xs_pad, masks = encoder_outs[0], encoder_outs[1] + + if layer_idx + 1 in self.interctc_layer_idx: + encoder_out = xs_pad + + # intermediate outputs are also normalized + if self.normalize_before: + encoder_out = self.after_norm(encoder_out) + + intermediate_outs.append((layer_idx + 1, encoder_out)) + + if self.interctc_use_conditioning: + ctc_out = ctc.softmax(encoder_out) + xs_pad = xs_pad + self.conditioning_layer(ctc_out) + + if self.normalize_before: + xs_pad = self.after_norm(xs_pad) + + olens = masks.squeeze(1).sum(1) + if len(intermediate_outs) > 0: + return (xs_pad, intermediate_outs), olens, None + return xs_pad, olens, None + + def _add_overlap_chunk(self, feats: np.ndarray, cache: dict = {}): + if len(cache) == 0: + return feats + cache["feats"] = to_device(cache["feats"], device=feats.device) + overlap_feats = torch.cat((cache["feats"], feats), dim=1) + cache["feats"] = overlap_feats[:, -(cache["chunk_size"][0] + cache["chunk_size"][2]) :, :] + return overlap_feats + + def forward_chunk( + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + cache: dict = None, + ctc: CTC = None, + ): + xs_pad *= self.output_size() ** 0.5 + if self.embed is None: + xs_pad = xs_pad + else: + xs_pad = self.embed(xs_pad, cache) + if cache["tail_chunk"]: + xs_pad = to_device(cache["feats"], device=xs_pad.device) + else: + xs_pad = self._add_overlap_chunk(xs_pad, cache) + encoder_outs = self.encoders0(xs_pad, None, None, None, None) + xs_pad, masks = encoder_outs[0], encoder_outs[1] + intermediate_outs = [] + if len(self.interctc_layer_idx) == 0: + encoder_outs = self.encoders(xs_pad, None, None, None, None) + xs_pad, masks = encoder_outs[0], encoder_outs[1] + else: + for layer_idx, encoder_layer in enumerate(self.encoders): + encoder_outs = encoder_layer(xs_pad, None, None, None, None) + xs_pad, masks = encoder_outs[0], encoder_outs[1] + if layer_idx + 1 in self.interctc_layer_idx: + encoder_out = xs_pad + + # intermediate outputs are also normalized + if self.normalize_before: + encoder_out = self.after_norm(encoder_out) + + intermediate_outs.append((layer_idx + 1, encoder_out)) + + if self.interctc_use_conditioning: + ctc_out = ctc.softmax(encoder_out) + xs_pad = xs_pad + self.conditioning_layer(ctc_out) + + if self.normalize_before: + xs_pad = self.after_norm(xs_pad) + + if len(intermediate_outs) > 0: + return (xs_pad, intermediate_outs), None, None + return xs_pad, ilens, None + + +class EncoderLayerSANMExport(nn.Module): + def __init__( + self, + model, + ): + """Construct an EncoderLayer object.""" + super().__init__() + self.self_attn = model.self_attn + self.feed_forward = model.feed_forward + self.norm1 = model.norm1 + self.norm2 = model.norm2 + self.in_size = model.in_size + self.size = model.size + + def forward(self, x, mask): + + residual = x + x = self.norm1(x) + x = self.self_attn(x, mask) + if self.in_size == self.size: + x = x + residual + residual = x + x = self.norm2(x) + x = self.feed_forward(x) + x = x + residual + + return x, mask + + +@tables.register("encoder_classes", "SANMEncoderChunkOptExport") +@tables.register("encoder_classes", "SANMEncoderExport") +class SANMEncoderExport(nn.Module): + def __init__( + self, + model, + max_seq_len=512, + feats_dim=560, + model_name="encoder", + onnx: bool = True, + ctc_linear: nn.Module = None, + ): + super().__init__() + self.embed = model.embed + if isinstance(self.embed, StreamSinusoidalPositionEncoder): + self.embed = None + self.model = model + self.feats_dim = feats_dim + self._output_size = model._output_size + + from funasr.utils.torch_function import sequence_mask + + self.make_pad_mask = sequence_mask(max_seq_len, flip=False) + + from funasr.models.sanm.attention import MultiHeadedAttentionSANMExport + + if hasattr(model, "encoders0"): + for i, d in enumerate(self.model.encoders0): + if isinstance(d.self_attn, MultiHeadedAttentionSANM): + d.self_attn = MultiHeadedAttentionSANMExport(d.self_attn) + self.model.encoders0[i] = EncoderLayerSANMExport(d) + + for i, d in enumerate(self.model.encoders): + if isinstance(d.self_attn, MultiHeadedAttentionSANM): + d.self_attn = MultiHeadedAttentionSANMExport(d.self_attn) + self.model.encoders[i] = EncoderLayerSANMExport(d) + + self.model_name = model_name + self.num_heads = model.encoders[0].self_attn.h + self.hidden_size = model.encoders[0].self_attn.linear_out.out_features + + self.ctc_linear = ctc_linear + + def prepare_mask(self, mask): + mask_3d_btd = mask[:, :, None] + if len(mask.shape) == 2: + mask_4d_bhlt = 1 - mask[:, None, None, :] + elif len(mask.shape) == 3: + mask_4d_bhlt = 1 - mask[:, None, :] + mask_4d_bhlt = mask_4d_bhlt * -10000.0 + + return mask_3d_btd, mask_4d_bhlt + + def forward(self, speech: torch.Tensor, speech_lengths: torch.Tensor, online: bool = False): + if not online: + speech = speech * self._output_size**0.5 + + mask = self.make_pad_mask(speech_lengths) + mask = self.prepare_mask(mask) + if self.embed is None: + xs_pad = speech + else: + xs_pad = self.embed(speech) + + encoder_outs = self.model.encoders0(xs_pad, mask) + xs_pad, masks = encoder_outs[0], encoder_outs[1] + + encoder_outs = self.model.encoders(xs_pad, mask) + xs_pad, masks = encoder_outs[0], encoder_outs[1] + + xs_pad = self.model.after_norm(xs_pad) + + if self.ctc_linear is not None: + xs_pad = self.ctc_linear(xs_pad) + xs_pad = F.softmax(xs_pad, dim=2) + + return xs_pad, speech_lengths + + def get_output_size(self): + return self.model.encoders[0].size + + def get_dummy_inputs(self): + feats = torch.randn(1, 100, self.feats_dim) + return feats + + def get_input_names(self): + return ["feats"] + + def get_output_names(self): + return ["encoder_out", "encoder_out_lens", "predictor_weight"] + + def get_dynamic_axes(self): + return { + "feats": {1: "feats_length"}, + "encoder_out": {1: "enc_out_length"}, + "predictor_weight": {1: "pre_out_length"}, + } diff --git a/funasr/models/sanm/model.py b/funasr/models/sanm/model.py new file mode 100644 index 0000000000000000000000000000000000000000..b3f84b2b9c50e7cb5f357e7b4f516d8fe07f470b --- /dev/null +++ b/funasr/models/sanm/model.py @@ -0,0 +1,28 @@ +#!/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 logging + +import torch + +from funasr.models.transformer.model import Transformer +from funasr.register import tables + + +@tables.register("model_classes", "SANM") +class SANM(Transformer): + """ + Author: Zhifu Gao, Shiliang Zhang, Ming Lei, Ian McLoughlin + San-m: Memory equipped self-attention for end-to-end speech recognition + https://arxiv.org/abs/2006.01713 + """ + + def __init__( + self, + *args, + **kwargs, + ): + + super().__init__(*args, **kwargs) diff --git a/funasr/models/sanm/multihead_att.py b/funasr/models/sanm/multihead_att.py new file mode 100644 index 0000000000000000000000000000000000000000..671d46011c87ae2c0635222226a936eab73c7581 --- /dev/null +++ b/funasr/models/sanm/multihead_att.py @@ -0,0 +1,241 @@ +import os +import math + +import torch +import torch.nn as nn + + +class MultiHeadedAttentionSANMExport(nn.Module): + def __init__(self, model): + super().__init__() + self.d_k = model.d_k + self.h = model.h + self.linear_out = model.linear_out + self.linear_q_k_v = model.linear_q_k_v + self.fsmn_block = model.fsmn_block + self.pad_fn = model.pad_fn + + self.attn = None + self.all_head_size = self.h * self.d_k + + def forward(self, x, mask): + mask_3d_btd, mask_4d_bhlt = mask + q_h, k_h, v_h, v = self.forward_qkv(x) + fsmn_memory = self.forward_fsmn(v, mask_3d_btd) + q_h = q_h * self.d_k ** (-0.5) + scores = torch.matmul(q_h, k_h.transpose(-2, -1)) + att_outs = self.forward_attention(v_h, scores, mask_4d_bhlt) + return att_outs + fsmn_memory + + def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor: + new_x_shape = x.size()[:-1] + (self.h, self.d_k) + x = x.view(new_x_shape) + return x.permute(0, 2, 1, 3) + + def forward_qkv(self, x): + q_k_v = self.linear_q_k_v(x) + q, k, v = torch.split(q_k_v, int(self.h * self.d_k), dim=-1) + q_h = self.transpose_for_scores(q) + k_h = self.transpose_for_scores(k) + v_h = self.transpose_for_scores(v) + return q_h, k_h, v_h, v + + def forward_fsmn(self, inputs, mask): + # b, t, d = inputs.size() + # mask = torch.reshape(mask, (b, -1, 1)) + inputs = inputs * mask + x = inputs.transpose(1, 2) + x = self.pad_fn(x) + x = self.fsmn_block(x) + x = x.transpose(1, 2) + x = x + inputs + x = x * mask + return x + + def forward_attention(self, value, scores, mask): + scores = scores + mask + + attn = torch.softmax(scores, dim=-1) + context_layer = torch.matmul(attn, value) # (batch, head, time1, d_k) + + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.view(new_context_layer_shape) + return self.linear_out(context_layer) # (batch, time1, d_model) + + +def preprocess_for_attn(x, mask, cache, pad_fn, kernel_size): + x = x * mask + x = x.transpose(1, 2) + if cache is None: + x = pad_fn(x) + else: + x = torch.cat((cache, x), dim=2) + cache = x[:, :, -(kernel_size - 1) :] + return x, cache + + +torch_version = tuple([int(i) for i in torch.__version__.split(".")[:2]]) +if torch_version >= (1, 8): + import torch.fx + + torch.fx.wrap("preprocess_for_attn") + + +class MultiHeadedAttentionSANMDecoderExport(nn.Module): + def __init__(self, model): + super().__init__() + self.fsmn_block = model.fsmn_block + self.pad_fn = model.pad_fn + self.kernel_size = model.kernel_size + self.attn = None + + def forward(self, inputs, mask, cache=None): + x, cache = preprocess_for_attn(inputs, mask, cache, self.pad_fn, self.kernel_size) + x = self.fsmn_block(x) + x = x.transpose(1, 2) + + x = x + inputs + x = x * mask + return x, cache + + +class MultiHeadedAttentionCrossAttExport(nn.Module): + def __init__(self, model): + super().__init__() + self.d_k = model.d_k + self.h = model.h + self.linear_q = model.linear_q + self.linear_k_v = model.linear_k_v + self.linear_out = model.linear_out + self.attn = None + self.all_head_size = self.h * self.d_k + + def forward(self, x, memory, memory_mask): + q, k, v = self.forward_qkv(x, memory) + scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) + return self.forward_attention(v, scores, memory_mask) + + def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor: + new_x_shape = x.size()[:-1] + (self.h, self.d_k) + x = x.view(new_x_shape) + return x.permute(0, 2, 1, 3) + + def forward_qkv(self, x, memory): + q = self.linear_q(x) + + k_v = self.linear_k_v(memory) + k, v = torch.split(k_v, int(self.h * self.d_k), dim=-1) + q = self.transpose_for_scores(q) + k = self.transpose_for_scores(k) + v = self.transpose_for_scores(v) + return q, k, v + + def forward_attention(self, value, scores, mask): + scores = scores + mask + + attn = torch.softmax(scores, dim=-1) + context_layer = torch.matmul(attn, value) # (batch, head, time1, d_k) + + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.view(new_context_layer_shape) + return self.linear_out(context_layer) # (batch, time1, d_model) + + +class OnnxMultiHeadedAttention(nn.Module): + def __init__(self, model): + super().__init__() + self.d_k = model.d_k + self.h = model.h + self.linear_q = model.linear_q + self.linear_k = model.linear_k + self.linear_v = model.linear_v + self.linear_out = model.linear_out + self.attn = None + self.all_head_size = self.h * self.d_k + + def forward(self, query, key, value, mask): + q, k, v = self.forward_qkv(query, key, value) + scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) + return self.forward_attention(v, scores, mask) + + def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor: + new_x_shape = x.size()[:-1] + (self.h, self.d_k) + x = x.view(new_x_shape) + return x.permute(0, 2, 1, 3) + + def forward_qkv(self, query, key, value): + q = self.linear_q(query) + k = self.linear_k(key) + v = self.linear_v(value) + q = self.transpose_for_scores(q) + k = self.transpose_for_scores(k) + v = self.transpose_for_scores(v) + return q, k, v + + def forward_attention(self, value, scores, mask): + scores = scores + mask + + attn = torch.softmax(scores, dim=-1) + context_layer = torch.matmul(attn, value) # (batch, head, time1, d_k) + + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.view(new_context_layer_shape) + return self.linear_out(context_layer) # (batch, time1, d_model) + + +class OnnxRelPosMultiHeadedAttention(OnnxMultiHeadedAttention): + def __init__(self, model): + super().__init__(model) + self.linear_pos = model.linear_pos + self.pos_bias_u = model.pos_bias_u + self.pos_bias_v = model.pos_bias_v + + def forward(self, query, key, value, pos_emb, mask): + q, k, v = self.forward_qkv(query, key, value) + q = q.transpose(1, 2) # (batch, time1, head, d_k) + + p = self.transpose_for_scores(self.linear_pos(pos_emb)) # (batch, head, time1, d_k) + + # (batch, head, time1, d_k) + q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2) + # (batch, head, time1, d_k) + q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2) + + # compute attention score + # first compute matrix a and matrix c + # as described in https://arxiv.org/abs/1901.02860 Section 3.3 + # (batch, head, time1, time2) + matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1)) + + # compute matrix b and matrix d + # (batch, head, time1, time1) + matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) + matrix_bd = self.rel_shift(matrix_bd) + + scores = (matrix_ac + matrix_bd) / math.sqrt(self.d_k) # (batch, head, time1, time2) + + return self.forward_attention(v, scores, mask) + + def rel_shift(self, x): + zero_pad = torch.zeros((*x.size()[:3], 1), device=x.device, dtype=x.dtype) + x_padded = torch.cat([zero_pad, x], dim=-1) + + x_padded = x_padded.view(*x.size()[:2], x.size(3) + 1, x.size(2)) + x = x_padded[:, :, 1:].view_as(x)[ + :, :, :, : x.size(-1) // 2 + 1 + ] # only keep the positions from 0 to time2 + return x + + def forward_attention(self, value, scores, mask): + scores = scores + mask + + attn = torch.softmax(scores, dim=-1) + context_layer = torch.matmul(attn, value) # (batch, head, time1, d_k) + + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.view(new_context_layer_shape) + return self.linear_out(context_layer) # (batch, time1, d_model) diff --git a/funasr/models/sanm/positionwise_feed_forward.py b/funasr/models/sanm/positionwise_feed_forward.py new file mode 100644 index 0000000000000000000000000000000000000000..e28ba8825e954157feea5eb9a9ad027fbb416e0b --- /dev/null +++ b/funasr/models/sanm/positionwise_feed_forward.py @@ -0,0 +1,33 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + + +"""Positionwise feed forward layer definition.""" + +import torch + +from funasr.models.transformer.layer_norm import LayerNorm + + +class PositionwiseFeedForwardDecoderSANM(torch.nn.Module): + """Positionwise feed forward layer. + + Args: + idim (int): Input dimenstion. + hidden_units (int): The number of hidden units. + dropout_rate (float): Dropout rate. + + """ + + def __init__(self, idim, hidden_units, dropout_rate, adim=None, activation=torch.nn.ReLU()): + """Construct an PositionwiseFeedForward object.""" + super(PositionwiseFeedForwardDecoderSANM, self).__init__() + self.w_1 = torch.nn.Linear(idim, hidden_units) + self.w_2 = torch.nn.Linear(hidden_units, idim if adim is None else adim, bias=False) + self.dropout = torch.nn.Dropout(dropout_rate) + self.activation = activation + self.norm = LayerNorm(hidden_units) + + def forward(self, x): + """Forward function.""" + return self.w_2(self.norm(self.dropout(self.activation(self.w_1(x))))) diff --git a/funasr/models/sanm/template.yaml b/funasr/models/sanm/template.yaml new file mode 100644 index 0000000000000000000000000000000000000000..316fe75cba9aa96400ffdabf76e4eb76287d31a9 --- /dev/null +++ b/funasr/models/sanm/template.yaml @@ -0,0 +1,121 @@ +# This is an example that demonstrates how to configure a model file. +# You can modify the configuration according to your own requirements. + +# to print the register_table: +# from funasr.register import tables +# tables.print() + +# network architecture +model: SANM +model_conf: + ctc_weight: 0.0 + lsm_weight: 0.1 + length_normalized_loss: true + +# encoder +encoder: SANMEncoder +encoder_conf: + output_size: 512 + attention_heads: 4 + linear_units: 2048 + num_blocks: 50 + dropout_rate: 0.1 + positional_dropout_rate: 0.1 + attention_dropout_rate: 0.1 + input_layer: pe + pos_enc_class: SinusoidalPositionEncoder + normalize_before: true + kernel_size: 11 + sanm_shfit: 0 + selfattention_layer_type: sanm + +# decoder +decoder: FsmnDecoder +decoder_conf: + attention_heads: 4 + linear_units: 2048 + num_blocks: 16 + dropout_rate: 0.1 + positional_dropout_rate: 0.1 + self_attention_dropout_rate: 0.1 + src_attention_dropout_rate: 0.1 + att_layer_num: 16 + kernel_size: 11 + sanm_shfit: 0 + + + +# frontend related +frontend: WavFrontend +frontend_conf: + fs: 16000 + window: hamming + n_mels: 80 + frame_length: 25 + frame_shift: 10 + lfr_m: 7 + lfr_n: 6 + +specaug: SpecAugLFR +specaug_conf: + apply_time_warp: false + time_warp_window: 5 + time_warp_mode: bicubic + apply_freq_mask: true + freq_mask_width_range: + - 0 + - 30 + lfr_rate: 6 + num_freq_mask: 1 + apply_time_mask: true + time_mask_width_range: + - 0 + - 12 + num_time_mask: 1 + +train_conf: + accum_grad: 1 + grad_clip: 5 + max_epoch: 150 + val_scheduler_criterion: + - valid + - acc + best_model_criterion: + - - valid + - acc + - max + keep_nbest_models: 10 + avg_nbest_model: 10 + log_interval: 50 + +optim: adam +optim_conf: + lr: 0.0005 +scheduler: warmuplr +scheduler_conf: + warmup_steps: 30000 + +dataset: AudioDataset +dataset_conf: + index_ds: IndexDSJsonl + batch_sampler: BatchSampler + batch_type: example # example or length + batch_size: 1 # if batch_type is example, batch_size is the numbers of samples; if length, batch_size is source_token_len+target_token_len; + max_token_length: 2048 # filter samples if source_token_len+target_token_len > max_token_length, + buffer_size: 500 + shuffle: True + num_workers: 0 + +tokenizer: CharTokenizer +tokenizer_conf: + unk_symbol: + split_with_space: true + + +ctc_conf: + dropout_rate: 0.0 + ctc_type: builtin + reduce: true + ignore_nan_grad: true + +normalize: null diff --git a/funasr/models/sanm_kws/__init__.py b/funasr/models/sanm_kws/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/models/sanm_kws/export_meta.py b/funasr/models/sanm_kws/export_meta.py new file mode 100644 index 0000000000000000000000000000000000000000..a91a22cb060d861adea62383d55ffe9bd87b8445 --- /dev/null +++ b/funasr/models/sanm_kws/export_meta.py @@ -0,0 +1,98 @@ +#!/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 types +import copy +import torch +from funasr.register import tables + + +def export_rebuild_model(model, **kwargs): + # self.device = kwargs.get("device") + is_onnx = kwargs.get("type", "onnx") == "onnx" + encoder_class = tables.encoder_classes.get(kwargs["encoder"] + "Export") + + if hasattr(model, "ctc"): + model.encoder = encoder_class( + model.encoder, + onnx=is_onnx, + feats_dim=kwargs.get("input_size", 560), + ctc_linear=model.ctc.ctc_lo + ) + else: + assert False, print(model) + model.encoder = encoder_class(model.encoder, onnx=is_onnx, feats_dim=kwargs.get("input_size", 560)) + + # from funasr.utils.torch_function import sequence_mask + # model.make_pad_mask = sequence_mask(max_seq_len=None, flip=False) + + encoder_model = copy.copy(model) + + # encoder + encoder_model.forward = types.MethodType(export_encoder_forward, encoder_model) + encoder_model.export_dummy_inputs = types.MethodType(export_encoder_dummy_inputs, encoder_model) + encoder_model.export_input_names = types.MethodType(export_encoder_input_names, encoder_model) + encoder_model.export_output_names = types.MethodType(export_encoder_output_names, encoder_model) + encoder_model.export_dynamic_axes = types.MethodType(export_encoder_dynamic_axes, encoder_model) + encoder_model.export_name = types.MethodType(export_encoder_name, encoder_model) + + return encoder_model + + +def export_encoder_forward( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, +): + # a. To device + batch = { + "speech": speech, + "speech_lengths": speech_lengths, + "online": True + } + # batch = to_device(batch, device=self.device) + + encoder_out, encoder_out_len = self.encoder(**batch) + # mask = self.make_pad_mask(encoder_out_len)[:, None, :] + # alphas, _ = self.predictor.forward_cnn(enc, mask) + + # return encoder_out, encoder_out_len, alphas + return encoder_out, encoder_out_len + + +def export_encoder_dummy_inputs(self): + speech = torch.randn(2, 30, 280) + speech_lengths = torch.tensor([6, 30], dtype=torch.int32) + return (speech, speech_lengths) + + +def export_encoder_input_names(self): + return ["speech", "speech_lengths"] + + +def export_encoder_output_names(self): + # return ["encoder_out", "encoder_out_len", "alphas"] + return ["encoder_out", "encoder_out_len"] + + +def export_encoder_dynamic_axes(self): + return { + "speech": { + 0: "batch_size", 1: "feats_length" + }, + "speech_lengths": { + 0: "batch_size", + }, + "encoder_out": { + 0: "batch_size", 1: "feats_length" + }, + "encoder_out_len": { + 0: "batch_size", + }, + } + + +def export_encoder_name(self): + return "encoder.onnx" diff --git a/funasr/models/sanm_kws/model.py b/funasr/models/sanm_kws/model.py new file mode 100644 index 0000000000000000000000000000000000000000..d8d8f955c4eb4e00f4bc7898500938699c0aa7c3 --- /dev/null +++ b/funasr/models/sanm_kws/model.py @@ -0,0 +1,266 @@ +#!/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.train_utils.device_funcs import to_device +from funasr.utils.datadir_writer import DatadirWriter +from funasr.models.paraformer.search import Hypothesis +from funasr.models.paraformer.cif_predictor import mae_loss +from funasr.train_utils.device_funcs import force_gatherable +from funasr.losses.label_smoothing_loss import LabelSmoothingLoss +from funasr.models.transformer.utils.add_sos_eos import add_sos_eos +from funasr.models.transformer.utils.nets_utils import make_pad_mask +from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank + + +@tables.register("model_classes", "SanmKWS") +class SanmKWS(torch.nn.Module): + """ + Author: Speech Lab of DAMO Academy, Alibaba Group + Paraformer: Fast and Accurate Parallel Transformer for Non-autoregressive End-to-End Speech Recognition + https://arxiv.org/abs/2206.08317 + """ + + def __init__( + self, + specaug: Optional[str] = None, + specaug_conf: Optional[Dict] = None, + normalize: str = None, + normalize_conf: Optional[Dict] = None, + encoder: str = None, + encoder_conf: Optional[Dict] = None, + ctc: str = None, + ctc_conf: Optional[Dict] = None, + ctc_weight: float = 1.0, + input_size: int = 360, + vocab_size: int = -1, + ignore_id: int = -1, + blank_id: int = 0, + sos: int = 1, + eos: int = 2, + **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() + + if ctc_conf is None: + ctc_conf = {} + ctc = CTC(odim=vocab_size, encoder_output_size=encoder_output_size, **ctc_conf) + + # note that eos is the same as sos (equivalent ID) + self.blank_id = blank_id + self.sos = sos if sos is not None else vocab_size - 1 + self.eos = eos if eos is not None else vocab_size - 1 + self.vocab_size = vocab_size + self.ignore_id = ignore_id + self.ctc_weight = ctc_weight + # self.token_list = token_list.copy() + # + # self.frontend = frontend + self.specaug = specaug + self.normalize = normalize + self.encoder = encoder + + self.ctc = ctc + self.error_calculator = None + + def forward( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + text: torch.Tensor, + text_lengths: torch.Tensor, + **kwargs, + ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]: + """Encoder + Decoder + Calc loss + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + text: (Batch, Length) + text_lengths: (Batch,) + """ + if len(text_lengths.size()) > 1: + text_lengths = text_lengths[:, 0] + if len(speech_lengths.size()) > 1: + speech_lengths = speech_lengths[:, 0] + batch_size = speech.shape[0] + + # Encoder + encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) + + # decoder: CTC branch + loss_ctc, cer_ctc = self._calc_ctc_loss( + encoder_out, encoder_out_lens, text, text_lengths + ) + + # Collect CTC branch stats + stats = dict() + stats["loss_ctc"] = loss_ctc.detach() if loss_ctc is not None else None + stats["cer_ctc"] = cer_ctc + stats["cer"] = cer_ctc + + loss = loss_ctc + + stats["loss"] = torch.clone(loss.detach()) + stats["batch_size"] = batch_size + + # force_gatherable: to-device and to-tensor if scalar for DataParallel + loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) + return loss, stats, weight + + def encode( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Encoder. Note that this method is used by asr_inference.py + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + ind: int + """ + with autocast(False): + # Data augmentation + if self.specaug is not None and self.training: + speech, speech_lengths = self.specaug(speech, speech_lengths) + + # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN + if self.normalize is not None: + speech, speech_lengths = self.normalize(speech, speech_lengths) + + # Forward encoder + encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths) + if isinstance(encoder_out, tuple): + encoder_out = encoder_out[0] + + return encoder_out, encoder_out_lens + + def _calc_ctc_loss( + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, + ): + # Calc CTC loss + loss_ctc = self.ctc(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens) + + # Calc CER using CTC + cer_ctc = None + if not self.training and self.error_calculator is not None: + ys_hat = self.ctc.argmax(encoder_out).data + cer_ctc = self.error_calculator(ys_hat.cpu(), ys_pad.cpu(), is_ctc=True) + return loss_ctc, cer_ctc + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): + keywords = kwargs.get("keywords") + from funasr.utils.kws_utils import KwsCtcPrefixDecoder + self.kws_decoder = KwsCtcPrefixDecoder( + ctc=self.ctc, + keywords=keywords, + token_list=tokenizer.token_list, + seg_dict=tokenizer.seg_dict, + ) + + meta_data = {} + if ( + isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank" + ): # fbank + speech, speech_lengths = data_in, data_lengths + if len(speech.shape) < 3: + speech = speech[None, :, :] + if speech_lengths is not None: + speech_lengths = speech_lengths.squeeze(-1) + else: + speech_lengths = speech.shape[1] + else: + # extract fbank feats + time1 = time.perf_counter() + audio_sample_list = load_audio_text_image_video( + data_in, + fs=frontend.fs, + audio_fs=kwargs.get("fs", 16000), + data_type=kwargs.get("data_type", "sound"), + tokenizer=tokenizer, + ) + time2 = time.perf_counter() + meta_data["load_data"] = f"{time2 - time1:0.3f}" + speech, speech_lengths = extract_fbank( + audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend + ) + time3 = time.perf_counter() + meta_data["extract_feat"] = f"{time3 - time2:0.3f}" + meta_data["batch_data_time"] = ( + speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 + ) + + speech = speech.to(device=kwargs["device"]) + speech_lengths = speech_lengths.to(device=kwargs["device"]) + # Encoder + if kwargs.get("fp16", False): + speech = speech.half() + encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) + if isinstance(encoder_out, tuple): + encoder_out = encoder_out[0] + + results = [] + if kwargs.get("output_dir") is not None: + if not hasattr(self, "writer"): + self.writer = DatadirWriter(kwargs.get("output_dir")) + + for i in range(encoder_out.size(0)): + x = encoder_out[i, : encoder_out_lens[i], :] + detect_result = self.kws_decoder.decode(x) + is_deted, det_keyword, det_score = detect_result[0], detect_result[1], detect_result[2] + + if is_deted: + self.writer["detect"][key[i]] = "detected " + det_keyword + " " + str(det_score) + det_info = "detected " + det_keyword + " " + str(det_score) + else: + self.writer["detect"][key[i]] = "rejected" + det_info = "rejected" + + result_i = {"key": key[i], "text": det_info} + results.append(result_i) + + return results, meta_data + + def export(self, **kwargs): + from .export_meta import export_rebuild_model + + if "max_seq_len" not in kwargs: + kwargs["max_seq_len"] = 512 + models = export_rebuild_model(model=self, **kwargs) + return models diff --git a/funasr/models/sanm_kws_streaming/__init__.py b/funasr/models/sanm_kws_streaming/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/models/sanm_kws_streaming/export_meta.py b/funasr/models/sanm_kws_streaming/export_meta.py new file mode 100644 index 0000000000000000000000000000000000000000..a91a22cb060d861adea62383d55ffe9bd87b8445 --- /dev/null +++ b/funasr/models/sanm_kws_streaming/export_meta.py @@ -0,0 +1,98 @@ +#!/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 types +import copy +import torch +from funasr.register import tables + + +def export_rebuild_model(model, **kwargs): + # self.device = kwargs.get("device") + is_onnx = kwargs.get("type", "onnx") == "onnx" + encoder_class = tables.encoder_classes.get(kwargs["encoder"] + "Export") + + if hasattr(model, "ctc"): + model.encoder = encoder_class( + model.encoder, + onnx=is_onnx, + feats_dim=kwargs.get("input_size", 560), + ctc_linear=model.ctc.ctc_lo + ) + else: + assert False, print(model) + model.encoder = encoder_class(model.encoder, onnx=is_onnx, feats_dim=kwargs.get("input_size", 560)) + + # from funasr.utils.torch_function import sequence_mask + # model.make_pad_mask = sequence_mask(max_seq_len=None, flip=False) + + encoder_model = copy.copy(model) + + # encoder + encoder_model.forward = types.MethodType(export_encoder_forward, encoder_model) + encoder_model.export_dummy_inputs = types.MethodType(export_encoder_dummy_inputs, encoder_model) + encoder_model.export_input_names = types.MethodType(export_encoder_input_names, encoder_model) + encoder_model.export_output_names = types.MethodType(export_encoder_output_names, encoder_model) + encoder_model.export_dynamic_axes = types.MethodType(export_encoder_dynamic_axes, encoder_model) + encoder_model.export_name = types.MethodType(export_encoder_name, encoder_model) + + return encoder_model + + +def export_encoder_forward( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, +): + # a. To device + batch = { + "speech": speech, + "speech_lengths": speech_lengths, + "online": True + } + # batch = to_device(batch, device=self.device) + + encoder_out, encoder_out_len = self.encoder(**batch) + # mask = self.make_pad_mask(encoder_out_len)[:, None, :] + # alphas, _ = self.predictor.forward_cnn(enc, mask) + + # return encoder_out, encoder_out_len, alphas + return encoder_out, encoder_out_len + + +def export_encoder_dummy_inputs(self): + speech = torch.randn(2, 30, 280) + speech_lengths = torch.tensor([6, 30], dtype=torch.int32) + return (speech, speech_lengths) + + +def export_encoder_input_names(self): + return ["speech", "speech_lengths"] + + +def export_encoder_output_names(self): + # return ["encoder_out", "encoder_out_len", "alphas"] + return ["encoder_out", "encoder_out_len"] + + +def export_encoder_dynamic_axes(self): + return { + "speech": { + 0: "batch_size", 1: "feats_length" + }, + "speech_lengths": { + 0: "batch_size", + }, + "encoder_out": { + 0: "batch_size", 1: "feats_length" + }, + "encoder_out_len": { + 0: "batch_size", + }, + } + + +def export_encoder_name(self): + return "encoder.onnx" diff --git a/funasr/models/sanm_kws_streaming/model.py b/funasr/models/sanm_kws_streaming/model.py new file mode 100644 index 0000000000000000000000000000000000000000..c459e7c4216e340f623b10cbeea85e0ccdb36224 --- /dev/null +++ b/funasr/models/sanm_kws_streaming/model.py @@ -0,0 +1,442 @@ +#!/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 typing import Dict, Tuple +from contextlib import contextmanager +from distutils.version import LooseVersion + +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.sanm_kws.model import SanmKWS +from funasr.models.paraformer.search import Hypothesis +from funasr.models.paraformer.cif_predictor import mae_loss +from funasr.train_utils.device_funcs import force_gatherable +from funasr.losses.label_smoothing_loss import LabelSmoothingLoss +from funasr.models.transformer.utils.add_sos_eos import add_sos_eos +from funasr.models.transformer.utils.nets_utils import make_pad_mask, pad_list +from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank + + +if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): + from torch.cuda.amp import autocast +else: + # Nothing to do if torch<1.6.0 + @contextmanager + def autocast(enabled=True): + yield + + +@tables.register("model_classes", "SanmKWSStreaming") +class SanmKWSStreaming(SanmKWS): + """ + Author: Speech Lab of DAMO Academy, Alibaba Group + Paraformer: Fast and Accurate Parallel Transformer for Non-autoregressive End-to-End Speech Recognition + https://arxiv.org/abs/2206.08317 + """ + + def __init__( + self, + *args, + **kwargs, + ): + super().__init__(*args, **kwargs) + + def forward( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + text: torch.Tensor, + text_lengths: torch.Tensor, + **kwargs, + ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]: + """Encoder + Decoder + Calc loss + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + text: (Batch, Length) + text_lengths: (Batch,) + """ + decoding_ind = kwargs.get("decoding_ind") + if len(text_lengths.size()) > 1: + text_lengths = text_lengths[:, 0] + if len(speech_lengths.size()) > 1: + speech_lengths = speech_lengths[:, 0] + + batch_size = speech.shape[0] + + # Encoder + if hasattr(self.encoder, "overlap_chunk_cls"): + ind = self.encoder.overlap_chunk_cls.random_choice(self.training, decoding_ind) + encoder_out, encoder_out_lens = self.encode(speech, speech_lengths, ind=ind) + else: + encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) + + # decoder: CTC branch + if hasattr(self.encoder, "overlap_chunk_cls"): + encoder_out_ctc, encoder_out_lens_ctc = self.encoder.overlap_chunk_cls.remove_chunk( + encoder_out, encoder_out_lens, chunk_outs=None + ) + else: + encoder_out_ctc, encoder_out_lens_ctc = encoder_out, encoder_out_lens + + loss_ctc, cer_ctc = self._calc_ctc_loss( + encoder_out_ctc, encoder_out_lens_ctc, text, text_lengths + ) + + # Collect CTC branch stats + stats = dict() + stats["loss_ctc"] = loss_ctc.detach() if loss_ctc is not None else None + stats["cer_ctc"] = cer_ctc + + loss = loss_ctc + + stats["cer"] = cer_ctc + stats["loss"] = torch.clone(loss.detach()) + + # force_gatherable: to-device and to-tensor if scalar for DataParallel + loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) + return loss, stats, weight + + def encode_chunk( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + cache: dict = None, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Frontend + Encoder. Note that this method is used by asr_inference.py + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + ind: int + """ + with autocast(False): + # Data augmentation + if self.specaug is not None and self.training: + speech, speech_lengths = self.specaug(speech, speech_lengths) + + # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN + if self.normalize is not None: + speech, speech_lengths = self.normalize(speech, speech_lengths) + + # Forward encoder + encoder_out, encoder_out_lens, _ = self.encoder.forward_chunk( + speech, speech_lengths, cache=cache["encoder"] + ) + + if isinstance(encoder_out, tuple): + encoder_out = encoder_out[0] + + return encoder_out, torch.tensor([encoder_out.size(1)]) + + def init_cache(self, cache: dict = {}, **kwargs): + chunk_size = kwargs.get("chunk_size", [0, 10, 5]) + encoder_chunk_look_back = kwargs.get("encoder_chunk_look_back", 0) + decoder_chunk_look_back = kwargs.get("decoder_chunk_look_back", 0) + batch_size = 1 + + enc_output_size = kwargs["encoder_conf"]["output_size"] + feats_dims = kwargs["frontend_conf"]["n_mels"] * kwargs["frontend_conf"]["lfr_m"] + cache_encoder = { + "start_idx": 0, + "cif_hidden": torch.zeros((batch_size, 1, enc_output_size)), + "cif_alphas": torch.zeros((batch_size, 1)), + "encoder_out": None, + "encoder_out_lens": None, + "chunk_size": chunk_size, + "encoder_chunk_look_back": encoder_chunk_look_back, + "last_chunk": False, + "opt": None, + "feats": torch.zeros((batch_size, chunk_size[0] + chunk_size[2], feats_dims)), + "tail_chunk": False, + } + cache["encoder"] = cache_encoder + + cache_decoder = { + "decode_fsmn": None, + "decoder_chunk_look_back": decoder_chunk_look_back, + "opt": None, + "chunk_size": chunk_size, + } + cache["decoder"] = cache_decoder + cache["frontend"] = {} + cache["prev_samples"] = torch.empty(0) + + return cache + + def generate_chunk( + self, + speech, + speech_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): + cache = kwargs.get("cache", {}) + speech = speech.to(device=kwargs["device"]) + speech_lengths = speech_lengths.to(device=kwargs["device"]) + + # Encoder + is_final = kwargs.get("is_final", False) + encoder_out, encoder_out_lens = self.encode_chunk( + speech, speech_lengths, cache=cache, is_final=is_final + ) + if isinstance(encoder_out, tuple): + encoder_out = encoder_out[0] + + chunk_size = cache["encoder"]["chunk_size"] + real_start_pos = chunk_size[0] + + if encoder_out_lens[0] > chunk_size[0] + chunk_size[1] + chunk_size[2]: + assert False, print("impossible case 1 !") + if encoder_out_lens[0] == chunk_size[0] + chunk_size[1] + chunk_size[2]: + real_end_pos = chunk_size[0] + chunk_size[1] + elif encoder_out_lens[0] > chunk_size[0] + chunk_size[1]: + real_end_pos = chunk_size[0] + chunk_size[1] + elif encoder_out_lens[0] > chunk_size[0]: + real_end_pos = encoder_out_lens[0] + else: + assert False, print("impossible case 2 !") + + encoder_out_accum = cache["encoder"]["encoder_out"] + if encoder_out_accum is not None: + encoder_out_accum = torch.cat((encoder_out_accum, encoder_out[:, real_start_pos:real_end_pos, :]), dim=1) + else: + encoder_out_accum = encoder_out[:, real_start_pos:real_end_pos, :] + cache["encoder"]["encoder_out"] = encoder_out_accum + + if cache["encoder"]["encoder_out_lens"] is not None: + cache["encoder"]["encoder_out_lens"][0] += real_end_pos - real_start_pos + else: + cache["encoder"]["encoder_out_lens"] = encoder_out_lens + cache["encoder"]["encoder_out_lens"][0] = real_end_pos - real_start_pos + + if is_final: + if kwargs.get("output_dir") is not None: + if not hasattr(self, "writer"): + self.writer = DatadirWriter(kwargs.get("output_dir")) + + results = [] + for i in range(encoder_out_accum.size(0)): + x = encoder_out_accum[i, : cache["encoder"]["encoder_out_lens"][i], :] + detect_result = self.kws_decoder.decode(x) + is_deted, det_keyword, det_score = detect_result[0], detect_result[1], detect_result[2] + + if is_deted: + self.writer["detect"][key[i]] = "detected " + det_keyword + " " + str(det_score) + det_info = "detected " + det_keyword + " " + str(det_score) + else: + self.writer["detect"][key[i]] = "rejected" + det_info = "rejected" + + result_i = {"key": key[i], "text": det_info} + results.append(result_i) + + return results + else: + return None + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + cache: dict = {}, + **kwargs, + ): + keywords = kwargs.get("keywords") + from funasr.utils.kws_utils import KwsCtcPrefixDecoder + self.kws_decoder = KwsCtcPrefixDecoder( + ctc=self.ctc, + keywords=keywords, + token_list=tokenizer.token_list, + seg_dict=tokenizer.seg_dict, + ) + + meta_data = {} + chunk_size = kwargs["chunk_size"] + chunk_stride_samples = int(chunk_size[1] * 960) # 600ms + first_chunk_padding_samples = int(chunk_size[2] * 960) # 600ms + + if len(cache) == 0: + self.init_cache(cache, **kwargs) + + time1 = time.perf_counter() + cfg = {"is_final": kwargs.get("is_final", False)} + 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, + cache=cfg, + ) + _is_final = cfg["is_final"] # if data_in is a file or url, set is_final=True + + time2 = time.perf_counter() + meta_data["load_data"] = f"{time2 - time1:0.3f}" + assert len(audio_sample_list) == 1, "batch_size must be set 1" + + audio_sample = torch.cat((cache["prev_samples"], audio_sample_list[0])) + + if len(audio_sample) < first_chunk_padding_samples: + print("key: {}, audio is too short for inference {}".format(key, len(audio_sample))) + + audio_sample_pre = audio_sample[0 : first_chunk_padding_samples] + feat_pre, feat_pre_lengths = extract_fbank( + [audio_sample_pre], + data_type=kwargs.get("data_type", "sound"), + frontend=frontend, + cache=cache["frontend"], + is_final=False, + ) + + audio_sample = audio_sample[first_chunk_padding_samples:] + audio_chunks = int(len(audio_sample) // chunk_stride_samples) + + for i in range(audio_chunks): + if i == 0: + cache["encoder"]["feats"][:, chunk_size[2] :, :] = feat_pre + + kwargs["is_final"] = False + audio_sample_i = audio_sample[i * chunk_stride_samples : (i + 1) * chunk_stride_samples] + + if kwargs["is_final"] and len(audio_sample_i) < 960: + cache["encoder"]["tail_chunk"] = True + speech = cache["encoder"]["feats"] + speech_lengths = torch.tensor([speech.shape[1]], dtype=torch.int64).to( + speech.device + ) + else: + # extract fbank feats + speech, speech_lengths = extract_fbank( + [audio_sample_i], + data_type=kwargs.get("data_type", "sound"), + frontend=frontend, + cache=cache["frontend"], + is_final=kwargs["is_final"], + ) + 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 + ) + + results_chunk_i = self.generate_chunk( + speech, + speech_lengths, + key=key, + tokenizer=tokenizer, + cache=cache, + frontend=frontend, + **kwargs, + ) + + # results_chunk_i must be None when is_final=False + assert results_chunk_i is None + + # process tail samples + tail_audio_sample = audio_sample[ audio_chunks * chunk_stride_samples: ] + if len(tail_audio_sample) < 960: + kwargs["is_final"] = _is_final + cache["encoder"]["tail_chunk"] = True + speech = cache["encoder"]["feats"] + speech_lengths = torch.tensor([speech.shape[1]], dtype=torch.int64).to( + speech.device + ) + results_chunk_tail = self.generate_chunk( + speech, + speech_lengths, + key=key, + tokenizer=tokenizer, + cache=cache, + frontend=frontend, + **kwargs, + ) + elif len(tail_audio_sample) <= first_chunk_padding_samples: + kwargs["is_final"] = _is_final + # extract fbank feats + # cache["encoder"]["tail_chunk"] = True # cannot be true + speech, speech_lengths = extract_fbank( + [ tail_audio_sample ], + data_type=kwargs.get("data_type", "sound"), + frontend=frontend, + cache=cache["frontend"], + is_final=kwargs["is_final"], + ) + results_chunk_tail = self.generate_chunk( + speech, + speech_lengths, + key=key, + tokenizer=tokenizer, + cache=cache, + frontend=frontend, + **kwargs, + ) + elif len(tail_audio_sample) > first_chunk_padding_samples and \ + len(tail_audio_sample) < chunk_stride_samples: + kwargs["is_final"] = False + # extract fbank feats + speech, speech_lengths = extract_fbank( + [ tail_audio_sample ], + data_type=kwargs.get("data_type", "sound"), + frontend=frontend, + cache=cache["frontend"], + is_final=kwargs["is_final"], + ) + results_chunk = self.generate_chunk( + speech, + speech_lengths, + key=key, + tokenizer=tokenizer, + cache=cache, + frontend=frontend, + **kwargs, + ) + # results_chunk must be None when is_final=False + assert results_chunk is None + + # push tail chunk + kwargs["is_final"] = _is_final + cache["encoder"]["tail_chunk"] = True + speech = cache["encoder"]["feats"] + speech_lengths = torch.tensor([speech.shape[1]], dtype=torch.int64).to( + speech.device + ) + results_chunk_tail = self.generate_chunk( + speech, + speech_lengths, + key=key, + tokenizer=tokenizer, + cache=cache, + frontend=frontend, + **kwargs, + ) + + result = results_chunk_tail + + if _is_final: + self.init_cache(cache, **kwargs) + + if kwargs.get("output_dir"): + if not hasattr(self, "writer"): + self.writer = DatadirWriter(kwargs.get("output_dir")) + + return result, meta_data + + def export(self, **kwargs): + from .export_meta import export_rebuild_model + + models = export_rebuild_model(model=self, **kwargs) + return models diff --git a/funasr/models/scama/__init__.py b/funasr/models/scama/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/models/scama/beam_search.py b/funasr/models/scama/beam_search.py new file mode 100644 index 0000000000000000000000000000000000000000..259f483b2c9949bd73bb810d6d238219eeba6320 --- /dev/null +++ b/funasr/models/scama/beam_search.py @@ -0,0 +1,978 @@ +"""Beam search module.""" + +from itertools import chain +import logging +from typing import Any +from typing import Dict +from typing import List +from typing import NamedTuple +from typing import Tuple +from typing import Union + +import torch + +from funasr.metrics.common import end_detect +from funasr.models.transformer.scorers.scorer_interface import PartialScorerInterface +from funasr.models.transformer.scorers.scorer_interface import ScorerInterface + + +class Hypothesis(NamedTuple): + """Hypothesis data type.""" + + yseq: torch.Tensor + score: Union[float, torch.Tensor] = 0 + scores: Dict[str, Union[float, torch.Tensor]] = dict() + states: Dict[str, Any] = dict() + + def asdict(self) -> dict: + """Convert data to JSON-friendly dict.""" + return self._replace( + yseq=self.yseq.tolist(), + score=float(self.score), + scores={k: float(v) for k, v in self.scores.items()}, + )._asdict() + + +class BeamSearchScama(torch.nn.Module): + """Beam search implementation.""" + + def __init__( + self, + scorers: Dict[str, ScorerInterface], + weights: Dict[str, float], + beam_size: int, + vocab_size: int, + sos: int, + eos: int, + token_list: List[str] = None, + pre_beam_ratio: float = 1.5, + pre_beam_score_key: str = None, + ): + """Initialize beam search. + + Args: + scorers (dict[str, ScorerInterface]): Dict of decoder modules + e.g., Decoder, CTCPrefixScorer, LM + The scorer will be ignored if it is `None` + weights (dict[str, float]): Dict of weights for each scorers + The scorer will be ignored if its weight is 0 + beam_size (int): The number of hypotheses kept during search + vocab_size (int): The number of vocabulary + sos (int): Start of sequence id + eos (int): End of sequence id + token_list (list[str]): List of tokens for debug log + pre_beam_score_key (str): key of scores to perform pre-beam search + pre_beam_ratio (float): beam size in the pre-beam search + will be `int(pre_beam_ratio * beam_size)` + + """ + super().__init__() + # set scorers + self.weights = weights + self.scorers = dict() + self.full_scorers = dict() + self.part_scorers = dict() + # this module dict is required for recursive cast + # `self.to(device, dtype)` in `recog.py` + self.nn_dict = torch.nn.ModuleDict() + for k, v in scorers.items(): + w = weights.get(k, 0) + if w == 0 or v is None: + continue + assert isinstance( + v, ScorerInterface + ), f"{k} ({type(v)}) does not implement ScorerInterface" + self.scorers[k] = v + if isinstance(v, PartialScorerInterface): + self.part_scorers[k] = v + else: + self.full_scorers[k] = v + if isinstance(v, torch.nn.Module): + self.nn_dict[k] = v + + # set configurations + self.sos = sos + self.eos = eos + self.token_list = token_list + self.pre_beam_size = int(pre_beam_ratio * beam_size) + self.beam_size = beam_size + self.n_vocab = vocab_size + if ( + pre_beam_score_key is not None + and pre_beam_score_key != "full" + and pre_beam_score_key not in self.full_scorers + ): + raise KeyError(f"{pre_beam_score_key} is not found in {self.full_scorers}") + self.pre_beam_score_key = pre_beam_score_key + self.do_pre_beam = ( + self.pre_beam_score_key is not None + and self.pre_beam_size < self.n_vocab + and len(self.part_scorers) > 0 + ) + + def init_hyp(self, x: torch.Tensor) -> List[Hypothesis]: + """Get an initial hypothesis data. + + Args: + x (torch.Tensor): The encoder output feature + + Returns: + Hypothesis: The initial hypothesis. + + """ + init_states = dict() + init_scores = dict() + for k, d in self.scorers.items(): + init_states[k] = d.init_state(x) + init_scores[k] = 0.0 + return [ + Hypothesis( + score=0.0, + scores=init_scores, + states=init_states, + yseq=torch.tensor([self.sos], device=x.device), + ) + ] + + @staticmethod + def append_token(xs: torch.Tensor, x: int) -> torch.Tensor: + """Append new token to prefix tokens. + + Args: + xs (torch.Tensor): The prefix token + x (int): The new token to append + + Returns: + torch.Tensor: New tensor contains: xs + [x] with xs.dtype and xs.device + + """ + x = torch.tensor([x], dtype=xs.dtype, device=xs.device) + return torch.cat((xs, x)) + + def score_full( + self, + hyp: Hypothesis, + x: torch.Tensor, + x_mask: torch.Tensor = None, + pre_acoustic_embeds: torch.Tensor = None, + ) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: + """Score new hypothesis by `self.full_scorers`. + + Args: + hyp (Hypothesis): Hypothesis with prefix tokens to score + x (torch.Tensor): Corresponding input feature + + Returns: + Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: Tuple of + score dict of `hyp` that has string keys of `self.full_scorers` + and tensor score values of shape: `(self.n_vocab,)`, + and state dict that has string keys + and state values of `self.full_scorers` + + """ + scores = dict() + states = dict() + for k, d in self.full_scorers.items(): + scores[k], states[k] = d.score( + hyp.yseq, hyp.states[k], x, x_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds + ) + return scores, states + + def score_partial( + self, hyp: Hypothesis, ids: torch.Tensor, x: torch.Tensor + ) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: + """Score new hypothesis by `self.part_scorers`. + + Args: + hyp (Hypothesis): Hypothesis with prefix tokens to score + ids (torch.Tensor): 1D tensor of new partial tokens to score + x (torch.Tensor): Corresponding input feature + + Returns: + Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: Tuple of + score dict of `hyp` that has string keys of `self.part_scorers` + and tensor score values of shape: `(len(ids),)`, + and state dict that has string keys + and state values of `self.part_scorers` + + """ + scores = dict() + states = dict() + for k, d in self.part_scorers.items(): + scores[k], states[k] = d.score_partial(hyp.yseq, ids, hyp.states[k], x) + return scores, states + + def beam( + self, weighted_scores: torch.Tensor, ids: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute topk full token ids and partial token ids. + + Args: + weighted_scores (torch.Tensor): The weighted sum scores for each tokens. + Its shape is `(self.n_vocab,)`. + ids (torch.Tensor): The partial token ids to compute topk + + Returns: + Tuple[torch.Tensor, torch.Tensor]: + The topk full token ids and partial token ids. + Their shapes are `(self.beam_size,)` + + """ + # no pre beam performed + if weighted_scores.size(0) == ids.size(0): + top_ids = weighted_scores.topk(self.beam_size)[1] + return top_ids, top_ids + + # mask pruned in pre-beam not to select in topk + tmp = weighted_scores[ids] + weighted_scores[:] = -float("inf") + weighted_scores[ids] = tmp + top_ids = weighted_scores.topk(self.beam_size)[1] + local_ids = weighted_scores[ids].topk(self.beam_size)[1] + return top_ids, local_ids + + @staticmethod + def merge_scores( + prev_scores: Dict[str, float], + next_full_scores: Dict[str, torch.Tensor], + full_idx: int, + next_part_scores: Dict[str, torch.Tensor], + part_idx: int, + ) -> Dict[str, torch.Tensor]: + """Merge scores for new hypothesis. + + Args: + prev_scores (Dict[str, float]): + The previous hypothesis scores by `self.scorers` + next_full_scores (Dict[str, torch.Tensor]): scores by `self.full_scorers` + full_idx (int): The next token id for `next_full_scores` + next_part_scores (Dict[str, torch.Tensor]): + scores of partial tokens by `self.part_scorers` + part_idx (int): The new token id for `next_part_scores` + + Returns: + Dict[str, torch.Tensor]: The new score dict. + Its keys are names of `self.full_scorers` and `self.part_scorers`. + Its values are scalar tensors by the scorers. + + """ + new_scores = dict() + for k, v in next_full_scores.items(): + new_scores[k] = prev_scores[k] + v[full_idx] + for k, v in next_part_scores.items(): + new_scores[k] = prev_scores[k] + v[part_idx] + return new_scores + + def merge_states(self, states: Any, part_states: Any, part_idx: int) -> Any: + """Merge states for new hypothesis. + + Args: + states: states of `self.full_scorers` + part_states: states of `self.part_scorers` + part_idx (int): The new token id for `part_scores` + + Returns: + Dict[str, torch.Tensor]: The new score dict. + Its keys are names of `self.full_scorers` and `self.part_scorers`. + Its values are states of the scorers. + + """ + new_states = dict() + for k, v in states.items(): + new_states[k] = v + for k, d in self.part_scorers.items(): + new_states[k] = d.select_state(part_states[k], part_idx) + return new_states + + def search( + self, + running_hyps: List[Hypothesis], + x: torch.Tensor, + x_mask: torch.Tensor = None, + pre_acoustic_embeds: torch.Tensor = None, + ) -> List[Hypothesis]: + """Search new tokens for running hypotheses and encoded speech x. + + Args: + running_hyps (List[Hypothesis]): Running hypotheses on beam + x (torch.Tensor): Encoded speech feature (T, D) + + Returns: + List[Hypotheses]: Best sorted hypotheses + + """ + best_hyps = [] + part_ids = torch.arange(self.n_vocab, device=x.device) # no pre-beam + for hyp in running_hyps: + # scoring + weighted_scores = torch.zeros(self.n_vocab, dtype=x.dtype, device=x.device) + scores, states = self.score_full( + hyp, x, x_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds + ) + for k in self.full_scorers: + weighted_scores += self.weights[k] * scores[k] + # partial scoring + if self.do_pre_beam: + pre_beam_scores = ( + weighted_scores + if self.pre_beam_score_key == "full" + else scores[self.pre_beam_score_key] + ) + part_ids = torch.topk(pre_beam_scores, self.pre_beam_size)[1] + part_scores, part_states = self.score_partial(hyp, part_ids, x) + for k in self.part_scorers: + weighted_scores[part_ids] += self.weights[k] * part_scores[k] + # add previous hyp score + weighted_scores += hyp.score + + # update hyps + for j, part_j in zip(*self.beam(weighted_scores, part_ids)): + # will be (2 x beam at most) + best_hyps.append( + Hypothesis( + score=weighted_scores[j], + yseq=self.append_token(hyp.yseq, j), + scores=self.merge_scores(hyp.scores, scores, j, part_scores, part_j), + states=self.merge_states(states, part_states, part_j), + ) + ) + + # sort and prune 2 x beam -> beam + best_hyps = sorted(best_hyps, key=lambda x: x.score, reverse=True)[ + : min(len(best_hyps), self.beam_size) + ] + return best_hyps + + def forward( + self, + x: torch.Tensor, + scama_mask: torch.Tensor = None, + pre_acoustic_embeds: torch.Tensor = None, + maxlenratio: float = 0.0, + minlenratio: float = 0.0, + maxlen: int = None, + minlen: int = 0, + ) -> List[Hypothesis]: + """Perform beam search. + + Args: + x (torch.Tensor): Encoded speech feature (T, D) + maxlenratio (float): Input length ratio to obtain max output length. + If maxlenratio=0.0 (default), it uses a end-detect function + to automatically find maximum hypothesis lengths + If maxlenratio<0.0, its absolute value is interpreted + as a constant max output length. + minlenratio (float): Input length ratio to obtain min output length. + + Returns: + list[Hypothesis]: N-best decoding results + + """ + if maxlen is None: + # set length bounds + if maxlenratio == 0: + maxlen = x.shape[0] + elif maxlenratio < 0: + maxlen = -1 * int(maxlenratio) + else: + maxlen = max(1, int(maxlenratio * x.size(0))) + minlen = int(minlenratio * x.size(0)) + + logging.info("decoder input length: " + str(x.shape[0])) + logging.info("max output length: " + str(maxlen)) + logging.info("min output length: " + str(minlen)) + + # main loop of prefix search + running_hyps = self.init_hyp(x) + ended_hyps = [] + for i in range(maxlen): + logging.debug("position " + str(i)) + mask_enc = None + if scama_mask is not None: + token_num_predictor = scama_mask.size(1) + token_id_slice = min(i, token_num_predictor - 1) + mask_enc = scama_mask[:, token_id_slice : token_id_slice + 1, :] + # if mask_enc.size(1) == 0: + # mask_enc = scama_mask[:, -2:-1, :] + # # mask_enc = torch.zeros_like(mask_enc) + pre_acoustic_embeds_cur = None + if pre_acoustic_embeds is not None: + b, t, d = pre_acoustic_embeds.size() + pad = torch.zeros((b, 1, d), dtype=pre_acoustic_embeds.dtype).to( + device=pre_acoustic_embeds.device + ) + pre_acoustic_embeds = torch.cat((pre_acoustic_embeds, pad), dim=1) + token_id_slice = min(i, t) + pre_acoustic_embeds_cur = pre_acoustic_embeds[ + :, token_id_slice : token_id_slice + 1, : + ] + + best = self.search( + running_hyps, x, x_mask=mask_enc, pre_acoustic_embeds=pre_acoustic_embeds_cur + ) + # post process of one iteration + running_hyps = self.post_process(i, maxlen, maxlenratio, best, ended_hyps) + # end detection + if maxlenratio == 0.0 and end_detect([h.asdict() for h in ended_hyps], i): + logging.info(f"end detected at {i}") + break + if len(running_hyps) == 0: + logging.info("no hypothesis. Finish decoding.") + break + else: + logging.debug(f"remained hypotheses: {len(running_hyps)}") + + nbest_hyps = sorted(ended_hyps, key=lambda x: x.score, reverse=True) + # check the number of hypotheses reaching to eos + if len(nbest_hyps) == 0: + logging.warning( + "there is no N-best results, perform recognition " "again with smaller minlenratio." + ) + return ( + [] + if minlenratio < 0.1 + else self.forward(x, maxlenratio, max(0.0, minlenratio - 0.1)) + ) + + # report the best result + for x in nbest_hyps: + yseq = "".join([self.token_list[x] for x in x.yseq]) + logging.debug("nbest: y: {}, yseq: {}, score: {}".format(x.yseq, yseq, x.score)) + best = nbest_hyps[0] + for k, v in best.scores.items(): + logging.info(f"{v:6.2f} * {self.weights[k]:3} = {v * self.weights[k]:6.2f} for {k}") + logging.info(f"total log probability: {best.score:.2f}") + logging.info(f"normalized log probability: {best.score / len(best.yseq):.2f}") + logging.info(f"total number of ended hypotheses: {len(nbest_hyps)}") + if self.token_list is not None: + logging.info( + "best hypo: " + "".join([self.token_list[x] for x in best.yseq[1:-1]]) + "\n" + ) + return nbest_hyps + + def post_process( + self, + i: int, + maxlen: int, + maxlenratio: float, + running_hyps: List[Hypothesis], + ended_hyps: List[Hypothesis], + ) -> List[Hypothesis]: + """Perform post-processing of beam search iterations. + + Args: + i (int): The length of hypothesis tokens. + maxlen (int): The maximum length of tokens in beam search. + maxlenratio (int): The maximum length ratio in beam search. + running_hyps (List[Hypothesis]): The running hypotheses in beam search. + ended_hyps (List[Hypothesis]): The ended hypotheses in beam search. + + Returns: + List[Hypothesis]: The new running hypotheses. + + """ + logging.debug(f"the number of running hypotheses: {len(running_hyps)}") + if self.token_list is not None: + logging.debug( + "best hypo: " + "".join([self.token_list[x] for x in running_hyps[0].yseq[1:]]) + ) + # add eos in the final loop to avoid that there are no ended hyps + if i == maxlen - 1: + logging.info("adding in the last position in the loop") + running_hyps = [ + h._replace(yseq=self.append_token(h.yseq, self.eos)) for h in running_hyps + ] + + # add ended hypotheses to a final list, and removed them from current hypotheses + # (this will be a problem, number of hyps < beam) + remained_hyps = [] + for hyp in running_hyps: + if hyp.yseq[-1] == self.eos: + # e.g., Word LM needs to add final score + for k, d in chain(self.full_scorers.items(), self.part_scorers.items()): + s = d.final_score(hyp.states[k]) + hyp.scores[k] += s + hyp = hyp._replace(score=hyp.score + self.weights[k] * s) + ended_hyps.append(hyp) + else: + remained_hyps.append(hyp) + return remained_hyps + + +class BeamSearchScamaStreaming(torch.nn.Module): + """Beam search implementation.""" + + def __init__( + self, + scorers: Dict[str, ScorerInterface], + weights: Dict[str, float], + beam_size: int, + vocab_size: int, + sos: int, + eos: int, + token_list: List[str] = None, + pre_beam_ratio: float = 1.5, + pre_beam_score_key: str = None, + ): + """Initialize beam search. + + Args: + scorers (dict[str, ScorerInterface]): Dict of decoder modules + e.g., Decoder, CTCPrefixScorer, LM + The scorer will be ignored if it is `None` + weights (dict[str, float]): Dict of weights for each scorers + The scorer will be ignored if its weight is 0 + beam_size (int): The number of hypotheses kept during search + vocab_size (int): The number of vocabulary + sos (int): Start of sequence id + eos (int): End of sequence id + token_list (list[str]): List of tokens for debug log + pre_beam_score_key (str): key of scores to perform pre-beam search + pre_beam_ratio (float): beam size in the pre-beam search + will be `int(pre_beam_ratio * beam_size)` + + """ + super().__init__() + # set scorers + self.weights = weights + self.scorers = dict() + self.full_scorers = dict() + self.part_scorers = dict() + # this module dict is required for recursive cast + # `self.to(device, dtype)` in `recog.py` + self.nn_dict = torch.nn.ModuleDict() + for k, v in scorers.items(): + w = weights.get(k, 0) + if w == 0 or v is None: + continue + assert isinstance( + v, ScorerInterface + ), f"{k} ({type(v)}) does not implement ScorerInterface" + self.scorers[k] = v + if isinstance(v, PartialScorerInterface): + self.part_scorers[k] = v + else: + self.full_scorers[k] = v + if isinstance(v, torch.nn.Module): + self.nn_dict[k] = v + + # set configurations + self.sos = sos + self.eos = eos + self.token_list = token_list + self.pre_beam_size = int(pre_beam_ratio * beam_size) + self.beam_size = beam_size + self.n_vocab = vocab_size + if ( + pre_beam_score_key is not None + and pre_beam_score_key != "full" + and pre_beam_score_key not in self.full_scorers + ): + raise KeyError(f"{pre_beam_score_key} is not found in {self.full_scorers}") + self.pre_beam_score_key = pre_beam_score_key + self.do_pre_beam = ( + self.pre_beam_score_key is not None + and self.pre_beam_size < self.n_vocab + and len(self.part_scorers) > 0 + ) + + def init_hyp(self, x) -> List[Hypothesis]: + """Get an initial hypothesis data. + + Args: + x (torch.Tensor): The encoder output feature + + Returns: + Hypothesis: The initial hypothesis. + + """ + init_states = dict() + init_scores = dict() + for k, d in self.scorers.items(): + init_states[k] = d.init_state(x) + init_scores[k] = 0.0 + return [ + Hypothesis( + score=0.0, + scores=init_scores, + states=init_states, + yseq=torch.tensor([self.sos], device=x.device), + ) + ] + + @staticmethod + def append_token(xs: torch.Tensor, x: int) -> torch.Tensor: + """Append new token to prefix tokens. + + Args: + xs (torch.Tensor): The prefix token + x (int): The new token to append + + Returns: + torch.Tensor: New tensor contains: xs + [x] with xs.dtype and xs.device + + """ + x = torch.tensor([x], dtype=xs.dtype, device=xs.device) + return torch.cat((xs, x)) + + def score_full( + self, + hyp: Hypothesis, + x: torch.Tensor, + x_mask: torch.Tensor = None, + pre_acoustic_embeds: torch.Tensor = None, + cache: dict = {}, + ) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: + """Score new hypothesis by `self.full_scorers`. + + Args: + hyp (Hypothesis): Hypothesis with prefix tokens to score + x (torch.Tensor): Corresponding input feature + + Returns: + Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: Tuple of + score dict of `hyp` that has string keys of `self.full_scorers` + and tensor score values of shape: `(self.n_vocab,)`, + and state dict that has string keys + and state values of `self.full_scorers` + + """ + scores = dict() + states = dict() + for k, d in self.full_scorers.items(): + scores[k], states[k] = d.score( + hyp.yseq, + hyp.states[k], + x, + x_mask=x_mask, + pre_acoustic_embeds=pre_acoustic_embeds, + cache=cache, + ) + return scores, states + + def score_partial( + self, hyp: Hypothesis, ids: torch.Tensor, x: torch.Tensor + ) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: + """Score new hypothesis by `self.part_scorers`. + + Args: + hyp (Hypothesis): Hypothesis with prefix tokens to score + ids (torch.Tensor): 1D tensor of new partial tokens to score + x (torch.Tensor): Corresponding input feature + + Returns: + Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: Tuple of + score dict of `hyp` that has string keys of `self.part_scorers` + and tensor score values of shape: `(len(ids),)`, + and state dict that has string keys + and state values of `self.part_scorers` + + """ + scores = dict() + states = dict() + for k, d in self.part_scorers.items(): + scores[k], states[k] = d.score_partial(hyp.yseq, ids, hyp.states[k], x) + return scores, states + + def beam( + self, weighted_scores: torch.Tensor, ids: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute topk full token ids and partial token ids. + + Args: + weighted_scores (torch.Tensor): The weighted sum scores for each tokens. + Its shape is `(self.n_vocab,)`. + ids (torch.Tensor): The partial token ids to compute topk + + Returns: + Tuple[torch.Tensor, torch.Tensor]: + The topk full token ids and partial token ids. + Their shapes are `(self.beam_size,)` + + """ + # no pre beam performed + if weighted_scores.size(0) == ids.size(0): + top_ids = weighted_scores.topk(self.beam_size)[1] + return top_ids, top_ids + + # mask pruned in pre-beam not to select in topk + tmp = weighted_scores[ids] + weighted_scores[:] = -float("inf") + weighted_scores[ids] = tmp + top_ids = weighted_scores.topk(self.beam_size)[1] + local_ids = weighted_scores[ids].topk(self.beam_size)[1] + return top_ids, local_ids + + @staticmethod + def merge_scores( + prev_scores: Dict[str, float], + next_full_scores: Dict[str, torch.Tensor], + full_idx: int, + next_part_scores: Dict[str, torch.Tensor], + part_idx: int, + ) -> Dict[str, torch.Tensor]: + """Merge scores for new hypothesis. + + Args: + prev_scores (Dict[str, float]): + The previous hypothesis scores by `self.scorers` + next_full_scores (Dict[str, torch.Tensor]): scores by `self.full_scorers` + full_idx (int): The next token id for `next_full_scores` + next_part_scores (Dict[str, torch.Tensor]): + scores of partial tokens by `self.part_scorers` + part_idx (int): The new token id for `next_part_scores` + + Returns: + Dict[str, torch.Tensor]: The new score dict. + Its keys are names of `self.full_scorers` and `self.part_scorers`. + Its values are scalar tensors by the scorers. + + """ + new_scores = dict() + for k, v in next_full_scores.items(): + new_scores[k] = prev_scores[k] + v[full_idx] + for k, v in next_part_scores.items(): + new_scores[k] = prev_scores[k] + v[part_idx] + return new_scores + + def merge_states(self, states: Any, part_states: Any, part_idx: int) -> Any: + """Merge states for new hypothesis. + + Args: + states: states of `self.full_scorers` + part_states: states of `self.part_scorers` + part_idx (int): The new token id for `part_scores` + + Returns: + Dict[str, torch.Tensor]: The new score dict. + Its keys are names of `self.full_scorers` and `self.part_scorers`. + Its values are states of the scorers. + + """ + new_states = dict() + for k, v in states.items(): + new_states[k] = v + for k, d in self.part_scorers.items(): + new_states[k] = d.select_state(part_states[k], part_idx) + return new_states + + def search( + self, + running_hyps: List[Hypothesis], + x: torch.Tensor, + x_mask: torch.Tensor = None, + pre_acoustic_embeds: torch.Tensor = None, + cache: dict = {}, + ) -> List[Hypothesis]: + """Search new tokens for running hypotheses and encoded speech x. + + Args: + running_hyps (List[Hypothesis]): Running hypotheses on beam + x (torch.Tensor): Encoded speech feature (T, D) + + Returns: + List[Hypotheses]: Best sorted hypotheses + + """ + best_hyps = [] + part_ids = torch.arange(self.n_vocab, device=x.device) # no pre-beam + for hyp in running_hyps: + # scoring + weighted_scores = torch.zeros(self.n_vocab, dtype=x.dtype, device=x.device) + scores, states = self.score_full( + hyp, x, x_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds, cache=cache + ) + for k in self.full_scorers: + weighted_scores += self.weights[k] * scores[k] + # partial scoring + if self.do_pre_beam: + pre_beam_scores = ( + weighted_scores + if self.pre_beam_score_key == "full" + else scores[self.pre_beam_score_key] + ) + part_ids = torch.topk(pre_beam_scores, self.pre_beam_size)[1] + part_scores, part_states = self.score_partial(hyp, part_ids, x) + for k in self.part_scorers: + weighted_scores[part_ids] += self.weights[k] * part_scores[k] + # add previous hyp score + weighted_scores += hyp.score + + # update hyps + for j, part_j in zip(*self.beam(weighted_scores, part_ids)): + # will be (2 x beam at most) + best_hyps.append( + Hypothesis( + score=weighted_scores[j], + yseq=self.append_token(hyp.yseq, j), + scores=self.merge_scores(hyp.scores, scores, j, part_scores, part_j), + states=self.merge_states(states, part_states, part_j), + ) + ) + + # sort and prune 2 x beam -> beam + best_hyps = sorted(best_hyps, key=lambda x: x.score, reverse=True)[ + : min(len(best_hyps), self.beam_size) + ] + return best_hyps + + def forward( + self, + x: torch.Tensor, + scama_mask: torch.Tensor = None, + pre_acoustic_embeds: torch.Tensor = None, + maxlenratio: float = 0.0, + minlenratio: float = 0.0, + maxlen: int = None, + minlen: int = 0, + cache: dict = {}, + ) -> List[Hypothesis]: + """Perform beam search. + + Args: + x (torch.Tensor): Encoded speech feature (T, D) + maxlenratio (float): Input length ratio to obtain max output length. + If maxlenratio=0.0 (default), it uses a end-detect function + to automatically find maximum hypothesis lengths + If maxlenratio<0.0, its absolute value is interpreted + as a constant max output length. + minlenratio (float): Input length ratio to obtain min output length. + + Returns: + list[Hypothesis]: N-best decoding results + + """ + if maxlen is None: + # set length bounds + if maxlenratio == 0: + maxlen = x.shape[0] + elif maxlenratio < 0: + maxlen = -1 * int(maxlenratio) + else: + maxlen = max(1, int(maxlenratio * x.size(0))) + minlen = int(minlenratio * x.size(0)) + + logging.info("decoder input length: " + str(x.shape[0])) + logging.info("max output length: " + str(maxlen)) + logging.info("min output length: " + str(minlen)) + + # main loop of prefix search + # running_hyps = self.init_hyp(x) + running_hyps = cache["running_hyps"] + ended_hyps = [] + for i in range(maxlen): + logging.debug("position " + str(i)) + mask_enc = None + # if scama_mask is not None: + # token_num_predictor = scama_mask.size(1) + # token_id_slice = min(i, token_num_predictor-1) + # mask_enc = scama_mask[:, token_id_slice:token_id_slice+1, :] + # # if mask_enc.size(1) == 0: + # # mask_enc = scama_mask[:, -2:-1, :] + # # # mask_enc = torch.zeros_like(mask_enc) + pre_acoustic_embeds_cur = None + if pre_acoustic_embeds is not None: + b, t, d = pre_acoustic_embeds.size() + pad = torch.zeros((b, 1, d), dtype=pre_acoustic_embeds.dtype).to( + device=pre_acoustic_embeds.device + ) + pre_acoustic_embeds = torch.cat((pre_acoustic_embeds, pad), dim=1) + token_id_slice = min(i, t) + pre_acoustic_embeds_cur = pre_acoustic_embeds[ + :, token_id_slice : token_id_slice + 1, : + ] + + best = self.search( + running_hyps, + x, + x_mask=mask_enc, + pre_acoustic_embeds=pre_acoustic_embeds_cur, + cache=cache["decoder"], + ) + # post process of one iteration + running_hyps = self.post_process(i, maxlen, maxlenratio, best, ended_hyps) + # end detection + if maxlenratio == 0.0 and end_detect([h.asdict() for h in ended_hyps], i): + logging.info(f"end detected at {i}") + break + if len(running_hyps) == 0: + logging.info("no hypothesis. Finish decoding.") + break + else: + logging.debug(f"remained hypotheses: {len(running_hyps)}") + + nbest_hyps = sorted(ended_hyps, key=lambda x: x.score, reverse=True) + # check the number of hypotheses reaching to eos + if len(nbest_hyps) == 0: + logging.warning( + "there is no N-best results, perform recognition " "again with smaller minlenratio." + ) + return ( + [] + if minlenratio < 0.1 + else self.forward(x, maxlenratio, max(0.0, minlenratio - 0.1)) + ) + + # report the best result + for x in nbest_hyps: + yseq = "".join([self.token_list[x] for x in x.yseq]) + logging.debug("nbest: y: {}, yseq: {}, score: {}".format(x.yseq, yseq, x.score)) + best = nbest_hyps[0] + for k, v in best.scores.items(): + logging.info(f"{v:6.2f} * {self.weights[k]:3} = {v * self.weights[k]:6.2f} for {k}") + logging.info(f"total log probability: {best.score:.2f}") + logging.info(f"normalized log probability: {best.score / len(best.yseq):.2f}") + logging.info(f"total number of ended hypotheses: {len(nbest_hyps)}") + if self.token_list is not None: + logging.info( + "best hypo: " + "".join([self.token_list[x] for x in best.yseq[1:-1]]) + "\n" + ) + return nbest_hyps + + def post_process( + self, + i: int, + maxlen: int, + maxlenratio: float, + running_hyps: List[Hypothesis], + ended_hyps: List[Hypothesis], + ) -> List[Hypothesis]: + """Perform post-processing of beam search iterations. + + Args: + i (int): The length of hypothesis tokens. + maxlen (int): The maximum length of tokens in beam search. + maxlenratio (int): The maximum length ratio in beam search. + running_hyps (List[Hypothesis]): The running hypotheses in beam search. + ended_hyps (List[Hypothesis]): The ended hypotheses in beam search. + + Returns: + List[Hypothesis]: The new running hypotheses. + + """ + logging.debug(f"the number of running hypotheses: {len(running_hyps)}") + if self.token_list is not None: + logging.debug( + "best hypo: " + "".join([self.token_list[x] for x in running_hyps[0].yseq[1:]]) + ) + # add eos in the final loop to avoid that there are no ended hyps + if i == maxlen - 1: + logging.info("adding in the last position in the loop") + running_hyps = [ + h._replace(yseq=self.append_token(h.yseq, self.eos)) for h in running_hyps + ] + + # add ended hypotheses to a final list, and removed them from current hypotheses + # (this will be a problem, number of hyps < beam) + remained_hyps = [] + for hyp in running_hyps: + if hyp.yseq[-1] == self.eos: + # e.g., Word LM needs to add final score + for k, d in chain(self.full_scorers.items(), self.part_scorers.items()): + s = d.final_score(hyp.states[k]) + hyp.scores[k] += s + hyp = hyp._replace(score=hyp.score + self.weights[k] * s) + ended_hyps.append(hyp) + else: + remained_hyps.append(hyp) + return remained_hyps diff --git a/funasr/models/scama/chunk_utilis.py b/funasr/models/scama/chunk_utilis.py new file mode 100644 index 0000000000000000000000000000000000000000..2fe3fa4e3879857a4b3ec2c8945101236dc5bfcd --- /dev/null +++ b/funasr/models/scama/chunk_utilis.py @@ -0,0 +1,493 @@ +import math +import torch +import numpy as np +import torch.nn.functional as F + +from funasr.models.scama.utils import sequence_mask +from funasr.models.transformer.utils.nets_utils import make_pad_mask + + +class overlap_chunk: + """ + Author: Speech Lab of DAMO Academy, Alibaba Group + San-m: Memory equipped self-attention for end-to-end speech recognition + https://arxiv.org/abs/2006.01713 + + """ + + def __init__( + self, + chunk_size: tuple = (16,), + stride: tuple = (10,), + pad_left: tuple = (0,), + encoder_att_look_back_factor: tuple = (1,), + shfit_fsmn: int = 0, + decoder_att_look_back_factor: tuple = (1,), + ): + + pad_left = self.check_chunk_size_args(chunk_size, pad_left) + encoder_att_look_back_factor = self.check_chunk_size_args( + chunk_size, encoder_att_look_back_factor + ) + decoder_att_look_back_factor = self.check_chunk_size_args( + chunk_size, decoder_att_look_back_factor + ) + ( + self.chunk_size, + self.stride, + self.pad_left, + self.encoder_att_look_back_factor, + self.decoder_att_look_back_factor, + ) = ( + chunk_size, + stride, + pad_left, + encoder_att_look_back_factor, + decoder_att_look_back_factor, + ) + self.shfit_fsmn = shfit_fsmn + self.x_add_mask = None + self.x_rm_mask = None + self.x_len = None + self.mask_shfit_chunk = None + self.mask_chunk_predictor = None + self.mask_att_chunk_encoder = None + self.mask_shift_att_chunk_decoder = None + self.chunk_outs = None + ( + self.chunk_size_cur, + self.stride_cur, + self.pad_left_cur, + self.encoder_att_look_back_factor_cur, + self.chunk_size_pad_shift_cur, + ) = (None, None, None, None, None) + + def check_chunk_size_args(self, chunk_size, x): + if len(x) < len(chunk_size): + x = [x[0] for i in chunk_size] + return x + + def get_chunk_size(self, ind: int = 0): + # with torch.no_grad: + chunk_size, stride, pad_left, encoder_att_look_back_factor, decoder_att_look_back_factor = ( + self.chunk_size[ind], + self.stride[ind], + self.pad_left[ind], + self.encoder_att_look_back_factor[ind], + self.decoder_att_look_back_factor[ind], + ) + ( + self.chunk_size_cur, + self.stride_cur, + self.pad_left_cur, + self.encoder_att_look_back_factor_cur, + self.chunk_size_pad_shift_cur, + self.decoder_att_look_back_factor_cur, + ) = ( + chunk_size, + stride, + pad_left, + encoder_att_look_back_factor, + chunk_size + self.shfit_fsmn, + decoder_att_look_back_factor, + ) + return ( + self.chunk_size_cur, + self.stride_cur, + self.pad_left_cur, + self.encoder_att_look_back_factor_cur, + self.chunk_size_pad_shift_cur, + ) + + def random_choice(self, training=True, decoding_ind=None): + chunk_num = len(self.chunk_size) + ind = 0 + if training and chunk_num > 1: + ind = torch.randint(0, chunk_num, ()).cpu().item() + if not training and decoding_ind is not None: + ind = int(decoding_ind) + + return ind + + def gen_chunk_mask(self, x_len, ind=0, num_units=1, num_units_predictor=1): + + with torch.no_grad(): + x_len = x_len.cpu().numpy() + x_len_max = x_len.max() + + chunk_size, stride, pad_left, encoder_att_look_back_factor, chunk_size_pad_shift = ( + self.get_chunk_size(ind) + ) + shfit_fsmn = self.shfit_fsmn + pad_right = chunk_size - stride - pad_left + + chunk_num_batch = np.ceil(x_len / stride).astype(np.int32) + x_len_chunk = ( + (chunk_num_batch - 1) * chunk_size_pad_shift + + shfit_fsmn + + pad_left + + 0 + + x_len + - (chunk_num_batch - 1) * stride + ) + x_len_chunk = x_len_chunk.astype(x_len.dtype) + x_len_chunk_max = x_len_chunk.max() + + chunk_num = int(math.ceil(x_len_max / stride)) + dtype = np.int32 + max_len_for_x_mask_tmp = max(chunk_size, x_len_max + pad_left) + x_add_mask = np.zeros([0, max_len_for_x_mask_tmp], dtype=dtype) + x_rm_mask = np.zeros([max_len_for_x_mask_tmp, 0], dtype=dtype) + mask_shfit_chunk = np.zeros([0, num_units], dtype=dtype) + mask_chunk_predictor = np.zeros([0, num_units_predictor], dtype=dtype) + mask_shift_att_chunk_decoder = np.zeros([0, 1], dtype=dtype) + mask_att_chunk_encoder = np.zeros([0, chunk_num * chunk_size_pad_shift], dtype=dtype) + for chunk_ids in range(chunk_num): + # x_mask add + fsmn_padding = np.zeros((shfit_fsmn, max_len_for_x_mask_tmp), dtype=dtype) + x_mask_cur = np.diag(np.ones(chunk_size, dtype=np.float32)) + x_mask_pad_left = np.zeros((chunk_size, chunk_ids * stride), dtype=dtype) + x_mask_pad_right = np.zeros((chunk_size, max_len_for_x_mask_tmp), dtype=dtype) + x_cur_pad = np.concatenate([x_mask_pad_left, x_mask_cur, x_mask_pad_right], axis=1) + x_cur_pad = x_cur_pad[:chunk_size, :max_len_for_x_mask_tmp] + x_add_mask_fsmn = np.concatenate([fsmn_padding, x_cur_pad], axis=0) + x_add_mask = np.concatenate([x_add_mask, x_add_mask_fsmn], axis=0) + + # x_mask rm + fsmn_padding = np.zeros((max_len_for_x_mask_tmp, shfit_fsmn), dtype=dtype) + padding_mask_left = np.zeros((max_len_for_x_mask_tmp, pad_left), dtype=dtype) + padding_mask_right = np.zeros((max_len_for_x_mask_tmp, pad_right), dtype=dtype) + x_mask_cur = np.diag(np.ones(stride, dtype=dtype)) + x_mask_cur_pad_top = np.zeros((chunk_ids * stride, stride), dtype=dtype) + x_mask_cur_pad_bottom = np.zeros((max_len_for_x_mask_tmp, stride), dtype=dtype) + x_rm_mask_cur = np.concatenate( + [x_mask_cur_pad_top, x_mask_cur, x_mask_cur_pad_bottom], axis=0 + ) + x_rm_mask_cur = x_rm_mask_cur[:max_len_for_x_mask_tmp, :stride] + x_rm_mask_cur_fsmn = np.concatenate( + [fsmn_padding, padding_mask_left, x_rm_mask_cur, padding_mask_right], axis=1 + ) + x_rm_mask = np.concatenate([x_rm_mask, x_rm_mask_cur_fsmn], axis=1) + + # fsmn_padding_mask + pad_shfit_mask = np.zeros([shfit_fsmn, num_units], dtype=dtype) + ones_1 = np.ones([chunk_size, num_units], dtype=dtype) + mask_shfit_chunk_cur = np.concatenate([pad_shfit_mask, ones_1], axis=0) + mask_shfit_chunk = np.concatenate([mask_shfit_chunk, mask_shfit_chunk_cur], axis=0) + + # predictor mask + zeros_1 = np.zeros([shfit_fsmn + pad_left, num_units_predictor], dtype=dtype) + ones_2 = np.ones([stride, num_units_predictor], dtype=dtype) + zeros_3 = np.zeros( + [chunk_size - stride - pad_left, num_units_predictor], dtype=dtype + ) + ones_zeros = np.concatenate([ones_2, zeros_3], axis=0) + mask_chunk_predictor_cur = np.concatenate([zeros_1, ones_zeros], axis=0) + mask_chunk_predictor = np.concatenate( + [mask_chunk_predictor, mask_chunk_predictor_cur], axis=0 + ) + + # encoder att mask + zeros_1_top = np.zeros([shfit_fsmn, chunk_num * chunk_size_pad_shift], dtype=dtype) + + zeros_2_num = max(chunk_ids - encoder_att_look_back_factor, 0) + zeros_2 = np.zeros([chunk_size, zeros_2_num * chunk_size_pad_shift], dtype=dtype) + + encoder_att_look_back_num = max(chunk_ids - zeros_2_num, 0) + zeros_2_left = np.zeros([chunk_size, shfit_fsmn], dtype=dtype) + ones_2_mid = np.ones([stride, stride], dtype=dtype) + zeros_2_bottom = np.zeros([chunk_size - stride, stride], dtype=dtype) + zeros_2_right = np.zeros([chunk_size, chunk_size - stride], dtype=dtype) + ones_2 = np.concatenate([ones_2_mid, zeros_2_bottom], axis=0) + ones_2 = np.concatenate([zeros_2_left, ones_2, zeros_2_right], axis=1) + ones_2 = np.tile(ones_2, [1, encoder_att_look_back_num]) + + zeros_3_left = np.zeros([chunk_size, shfit_fsmn], dtype=dtype) + ones_3_right = np.ones([chunk_size, chunk_size], dtype=dtype) + ones_3 = np.concatenate([zeros_3_left, ones_3_right], axis=1) + + zeros_remain_num = max(chunk_num - 1 - chunk_ids, 0) + zeros_remain = np.zeros( + [chunk_size, zeros_remain_num * chunk_size_pad_shift], dtype=dtype + ) + + ones2_bottom = np.concatenate([zeros_2, ones_2, ones_3, zeros_remain], axis=1) + mask_att_chunk_encoder_cur = np.concatenate([zeros_1_top, ones2_bottom], axis=0) + mask_att_chunk_encoder = np.concatenate( + [mask_att_chunk_encoder, mask_att_chunk_encoder_cur], axis=0 + ) + + # decoder fsmn_shift_att_mask + zeros_1 = np.zeros([shfit_fsmn, 1]) + ones_1 = np.ones([chunk_size, 1]) + mask_shift_att_chunk_decoder_cur = np.concatenate([zeros_1, ones_1], axis=0) + mask_shift_att_chunk_decoder = np.concatenate( + [mask_shift_att_chunk_decoder, mask_shift_att_chunk_decoder_cur], axis=0 + ) + + self.x_add_mask = x_add_mask[:x_len_chunk_max, : x_len_max + pad_left] + self.x_len_chunk = x_len_chunk + self.x_rm_mask = x_rm_mask[:x_len_max, :x_len_chunk_max] + self.x_len = x_len + self.mask_shfit_chunk = mask_shfit_chunk[:x_len_chunk_max, :] + self.mask_chunk_predictor = mask_chunk_predictor[:x_len_chunk_max, :] + self.mask_att_chunk_encoder = mask_att_chunk_encoder[:x_len_chunk_max, :x_len_chunk_max] + self.mask_shift_att_chunk_decoder = mask_shift_att_chunk_decoder[:x_len_chunk_max, :] + self.chunk_outs = ( + self.x_add_mask, + self.x_len_chunk, + self.x_rm_mask, + self.x_len, + self.mask_shfit_chunk, + self.mask_chunk_predictor, + self.mask_att_chunk_encoder, + self.mask_shift_att_chunk_decoder, + ) + + return self.chunk_outs + + def split_chunk(self, x, x_len, chunk_outs): + """ + :param x: (b, t, d) + :param x_length: (b) + :param ind: int + :return: + """ + x = x[:, : x_len.max(), :] + b, t, d = x.size() + x_len_mask = (~make_pad_mask(x_len, maxlen=t)).to(x.device) + x *= x_len_mask[:, :, None] + + x_add_mask = self.get_x_add_mask(chunk_outs, x.device, dtype=x.dtype) + x_len_chunk = self.get_x_len_chunk(chunk_outs, x_len.device, dtype=x_len.dtype) + pad = (0, 0, self.pad_left_cur, 0) + x = F.pad(x, pad, "constant", 0.0) + b, t, d = x.size() + x = torch.transpose(x, 1, 0) + x = torch.reshape(x, [t, -1]) + x_chunk = torch.mm(x_add_mask, x) + x_chunk = torch.reshape(x_chunk, [-1, b, d]).transpose(1, 0) + + return x_chunk, x_len_chunk + + def remove_chunk(self, x_chunk, x_len_chunk, chunk_outs): + x_chunk = x_chunk[:, : x_len_chunk.max(), :] + b, t, d = x_chunk.size() + x_len_chunk_mask = (~make_pad_mask(x_len_chunk, maxlen=t)).to(x_chunk.device) + x_chunk *= x_len_chunk_mask[:, :, None] + + x_rm_mask = self.get_x_rm_mask(chunk_outs, x_chunk.device, dtype=x_chunk.dtype) + x_len = self.get_x_len(chunk_outs, x_len_chunk.device, dtype=x_len_chunk.dtype) + x_chunk = torch.transpose(x_chunk, 1, 0) + x_chunk = torch.reshape(x_chunk, [t, -1]) + x = torch.mm(x_rm_mask, x_chunk) + x = torch.reshape(x, [-1, b, d]).transpose(1, 0) + + return x, x_len + + def get_x_add_mask(self, chunk_outs=None, device="cpu", idx=0, dtype=torch.float32): + with torch.no_grad(): + x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] + x = torch.from_numpy(x).type(dtype).to(device) + return x + + def get_x_len_chunk(self, chunk_outs=None, device="cpu", idx=1, dtype=torch.float32): + with torch.no_grad(): + x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] + x = torch.from_numpy(x).type(dtype).to(device) + return x + + def get_x_rm_mask(self, chunk_outs=None, device="cpu", idx=2, dtype=torch.float32): + with torch.no_grad(): + x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] + x = torch.from_numpy(x).type(dtype).to(device) + return x + + def get_x_len(self, chunk_outs=None, device="cpu", idx=3, dtype=torch.float32): + with torch.no_grad(): + x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] + x = torch.from_numpy(x).type(dtype).to(device) + return x + + def get_mask_shfit_chunk( + self, chunk_outs=None, device="cpu", batch_size=1, num_units=1, idx=4, dtype=torch.float32 + ): + with torch.no_grad(): + x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] + x = np.tile( + x[ + None, + :, + :, + ], + [batch_size, 1, num_units], + ) + x = torch.from_numpy(x).type(dtype).to(device) + return x + + def get_mask_chunk_predictor( + self, chunk_outs=None, device="cpu", batch_size=1, num_units=1, idx=5, dtype=torch.float32 + ): + with torch.no_grad(): + x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] + x = np.tile( + x[ + None, + :, + :, + ], + [batch_size, 1, num_units], + ) + x = torch.from_numpy(x).type(dtype).to(device) + return x + + def get_mask_att_chunk_encoder( + self, chunk_outs=None, device="cpu", batch_size=1, idx=6, dtype=torch.float32 + ): + with torch.no_grad(): + x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] + x = np.tile( + x[ + None, + :, + :, + ], + [batch_size, 1, 1], + ) + x = torch.from_numpy(x).type(dtype).to(device) + return x + + def get_mask_shift_att_chunk_decoder( + self, chunk_outs=None, device="cpu", batch_size=1, idx=7, dtype=torch.float32 + ): + with torch.no_grad(): + x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] + x = np.tile(x[None, None, :, 0], [batch_size, 1, 1]) + x = torch.from_numpy(x).type(dtype).to(device) + return x + + +def build_scama_mask_for_cross_attention_decoder( + predictor_alignments: torch.Tensor, + encoder_sequence_length: torch.Tensor, + chunk_size: int = 5, + encoder_chunk_size: int = 5, + attention_chunk_center_bias: int = 0, + attention_chunk_size: int = 1, + attention_chunk_type: str = "chunk", + step=None, + predictor_mask_chunk_hopping: torch.Tensor = None, + decoder_att_look_back_factor: int = 1, + mask_shift_att_chunk_decoder: torch.Tensor = None, + target_length: torch.Tensor = None, + is_training=True, + dtype: torch.dtype = torch.float32, +): + with torch.no_grad(): + device = predictor_alignments.device + batch_size, chunk_num = predictor_alignments.size() + maximum_encoder_length = encoder_sequence_length.max().item() + int_type = predictor_alignments.dtype + if not is_training: + target_length = predictor_alignments.sum(dim=-1).type(encoder_sequence_length.dtype) + maximum_target_length = target_length.max() + predictor_alignments_cumsum = torch.cumsum(predictor_alignments, dim=1) + predictor_alignments_cumsum = predictor_alignments_cumsum[:, None, :].repeat( + 1, maximum_target_length, 1 + ) + + index = torch.ones([batch_size, maximum_target_length], dtype=int_type).to(device) + index = torch.cumsum(index, dim=1) + index = index[:, :, None].repeat(1, 1, chunk_num) + + index_div = torch.floor(torch.divide(predictor_alignments_cumsum, index)).type(int_type) + index_div_bool_zeros = index_div == 0 + index_div_bool_zeros_count = torch.sum(index_div_bool_zeros.type(int_type), dim=-1) + 1 + + index_div_bool_zeros_count = torch.clip(index_div_bool_zeros_count, min=1, max=chunk_num) + + index_div_bool_zeros_count *= chunk_size + index_div_bool_zeros_count += attention_chunk_center_bias + index_div_bool_zeros_count = torch.clip( + index_div_bool_zeros_count - 1, min=0, max=maximum_encoder_length + ) + index_div_bool_zeros_count_ori = index_div_bool_zeros_count + + index_div_bool_zeros_count = ( + torch.floor(index_div_bool_zeros_count / encoder_chunk_size) + 1 + ) * encoder_chunk_size + max_len_chunk = math.ceil(maximum_encoder_length / encoder_chunk_size) * encoder_chunk_size + + mask_flip, mask_flip2 = None, None + if attention_chunk_size is not None: + index_div_bool_zeros_count_beg = index_div_bool_zeros_count - attention_chunk_size + index_div_bool_zeros_count_beg = torch.clip( + index_div_bool_zeros_count_beg, 0, max_len_chunk + ) + index_div_bool_zeros_count_beg_mask = sequence_mask( + index_div_bool_zeros_count_beg, maxlen=max_len_chunk, dtype=int_type, device=device + ) + mask_flip = 1 - index_div_bool_zeros_count_beg_mask + attention_chunk_size2 = attention_chunk_size * (decoder_att_look_back_factor + 1) + index_div_bool_zeros_count_beg = index_div_bool_zeros_count - attention_chunk_size2 + + index_div_bool_zeros_count_beg = torch.clip( + index_div_bool_zeros_count_beg, 0, max_len_chunk + ) + index_div_bool_zeros_count_beg_mask = sequence_mask( + index_div_bool_zeros_count_beg, maxlen=max_len_chunk, dtype=int_type, device=device + ) + mask_flip2 = 1 - index_div_bool_zeros_count_beg_mask + + mask = sequence_mask( + index_div_bool_zeros_count, maxlen=max_len_chunk, dtype=dtype, device=device + ) + + if predictor_mask_chunk_hopping is not None: + b, k, t = mask.size() + predictor_mask_chunk_hopping = predictor_mask_chunk_hopping[:, None, :, 0].repeat( + 1, k, 1 + ) + + mask_mask_flip = mask + if mask_flip is not None: + mask_mask_flip = mask_flip * mask + + def _fn(): + mask_sliced = mask[:b, :k, encoder_chunk_size:t] + zero_pad_right = torch.zeros( + [b, k, encoder_chunk_size], dtype=mask_sliced.dtype + ).to(device) + mask_sliced = torch.cat([mask_sliced, zero_pad_right], dim=2) + _, _, tt = predictor_mask_chunk_hopping.size() + pad_right_p = max_len_chunk - tt + predictor_mask_chunk_hopping_pad = torch.nn.functional.pad( + predictor_mask_chunk_hopping, [0, pad_right_p], "constant", 0 + ) + masked = mask_sliced * predictor_mask_chunk_hopping_pad + + mask_true = mask_mask_flip + masked + return mask_true + + mask = _fn() if t > chunk_size else mask_mask_flip + + if mask_flip2 is not None: + mask *= mask_flip2 + + mask_target = sequence_mask( + target_length, maxlen=maximum_target_length, dtype=mask.dtype, device=device + ) + mask = mask[:, :maximum_target_length, :] * mask_target[:, :, None] + + mask_len = sequence_mask( + encoder_sequence_length, maxlen=maximum_encoder_length, dtype=mask.dtype, device=device + ) + mask = mask[:, :, :maximum_encoder_length] * mask_len[:, None, :] + + if attention_chunk_type == "full": + mask = torch.ones_like(mask).to(device) + if mask_shift_att_chunk_decoder is not None: + mask = mask * mask_shift_att_chunk_decoder + mask = mask[:, :maximum_target_length, :maximum_encoder_length].type(dtype).to(device) + + return mask diff --git a/funasr/models/scama/decoder.py b/funasr/models/scama/decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..31b235778ebadcd74a1e9d99ee11a94f68991246 --- /dev/null +++ b/funasr/models/scama/decoder.py @@ -0,0 +1,498 @@ +#!/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) + +from typing import List +from typing import Tuple +import logging +import torch +import torch.nn as nn +import numpy as np + +from funasr.models.scama import utils as myutils +from funasr.models.transformer.decoder import BaseTransformerDecoder + +from funasr.models.sanm.attention import ( + MultiHeadedAttentionSANMDecoder, + MultiHeadedAttentionCrossAtt, +) +from funasr.models.transformer.embedding import PositionalEncoding +from funasr.models.transformer.layer_norm import LayerNorm +from funasr.models.sanm.positionwise_feed_forward import PositionwiseFeedForwardDecoderSANM +from funasr.models.transformer.utils.repeat import repeat + +from funasr.register import tables + + +class DecoderLayerSANM(nn.Module): + """Single decoder layer module. + + Args: + size (int): Input dimension. + self_attn (torch.nn.Module): Self-attention module instance. + `MultiHeadedAttention` instance can be used as the argument. + src_attn (torch.nn.Module): Self-attention module instance. + `MultiHeadedAttention` instance can be used as the argument. + feed_forward (torch.nn.Module): Feed-forward module instance. + `PositionwiseFeedForward`, `MultiLayeredConv1d`, or `Conv1dLinear` instance + can be used as the argument. + dropout_rate (float): Dropout rate. + normalize_before (bool): Whether to use layer_norm before the first block. + concat_after (bool): Whether to concat attention layer's input and output. + if True, additional linear will be applied. + i.e. x -> x + linear(concat(x, att(x))) + if False, no additional linear will be applied. i.e. x -> x + att(x) + + + """ + + def __init__( + self, + size, + self_attn, + src_attn, + feed_forward, + dropout_rate, + normalize_before=True, + concat_after=False, + ): + """Construct an DecoderLayer object.""" + super(DecoderLayerSANM, self).__init__() + self.size = size + self.self_attn = self_attn + self.src_attn = src_attn + self.feed_forward = feed_forward + self.norm1 = LayerNorm(size) + if self_attn is not None: + self.norm2 = LayerNorm(size) + if src_attn is not None: + self.norm3 = LayerNorm(size) + self.dropout = nn.Dropout(dropout_rate) + self.normalize_before = normalize_before + self.concat_after = concat_after + if self.concat_after: + self.concat_linear1 = nn.Linear(size + size, size) + self.concat_linear2 = nn.Linear(size + size, size) + + def forward(self, tgt, tgt_mask, memory, memory_mask=None, cache=None): + """Compute decoded features. + + Args: + tgt (torch.Tensor): Input tensor (#batch, maxlen_out, size). + tgt_mask (torch.Tensor): Mask for input tensor (#batch, maxlen_out). + memory (torch.Tensor): Encoded memory, float32 (#batch, maxlen_in, size). + memory_mask (torch.Tensor): Encoded memory mask (#batch, maxlen_in). + cache (List[torch.Tensor]): List of cached tensors. + Each tensor shape should be (#batch, maxlen_out - 1, size). + + Returns: + torch.Tensor: Output tensor(#batch, maxlen_out, size). + torch.Tensor: Mask for output tensor (#batch, maxlen_out). + torch.Tensor: Encoded memory (#batch, maxlen_in, size). + torch.Tensor: Encoded memory mask (#batch, maxlen_in). + + """ + # tgt = self.dropout(tgt) + residual = tgt + if self.normalize_before: + tgt = self.norm1(tgt) + tgt = self.feed_forward(tgt) + + x = tgt + if self.self_attn: + if self.normalize_before: + tgt = self.norm2(tgt) + x, _ = self.self_attn(tgt, tgt_mask) + x = residual + self.dropout(x) + + if self.src_attn is not None: + residual = x + if self.normalize_before: + x = self.norm3(x) + + x = residual + self.dropout(self.src_attn(x, memory, memory_mask)) + + return x, tgt_mask, memory, memory_mask, cache + + def forward_one_step(self, tgt, tgt_mask, memory, memory_mask=None, cache=None): + """Compute decoded features. + + Args: + tgt (torch.Tensor): Input tensor (#batch, maxlen_out, size). + tgt_mask (torch.Tensor): Mask for input tensor (#batch, maxlen_out). + memory (torch.Tensor): Encoded memory, float32 (#batch, maxlen_in, size). + memory_mask (torch.Tensor): Encoded memory mask (#batch, maxlen_in). + cache (List[torch.Tensor]): List of cached tensors. + Each tensor shape should be (#batch, maxlen_out - 1, size). + + Returns: + torch.Tensor: Output tensor(#batch, maxlen_out, size). + torch.Tensor: Mask for output tensor (#batch, maxlen_out). + torch.Tensor: Encoded memory (#batch, maxlen_in, size). + torch.Tensor: Encoded memory mask (#batch, maxlen_in). + + """ + # tgt = self.dropout(tgt) + residual = tgt + if self.normalize_before: + tgt = self.norm1(tgt) + tgt = self.feed_forward(tgt) + + x = tgt + if self.self_attn: + if self.normalize_before: + tgt = self.norm2(tgt) + if self.training: + cache = None + x, cache = self.self_attn(tgt, tgt_mask, cache=cache) + x = residual + self.dropout(x) + + if self.src_attn is not None: + residual = x + if self.normalize_before: + x = self.norm3(x) + + x = residual + self.dropout(self.src_attn(x, memory, memory_mask)) + + return x, tgt_mask, memory, memory_mask, cache + + def forward_chunk( + self, tgt, memory, fsmn_cache=None, opt_cache=None, chunk_size=None, look_back=0 + ): + """Compute decoded features. + + Args: + tgt (torch.Tensor): Input tensor (#batch, maxlen_out, size). + tgt_mask (torch.Tensor): Mask for input tensor (#batch, maxlen_out). + memory (torch.Tensor): Encoded memory, float32 (#batch, maxlen_in, size). + memory_mask (torch.Tensor): Encoded memory mask (#batch, maxlen_in). + cache (List[torch.Tensor]): List of cached tensors. + Each tensor shape should be (#batch, maxlen_out - 1, size). + + Returns: + torch.Tensor: Output tensor(#batch, maxlen_out, size). + torch.Tensor: Mask for output tensor (#batch, maxlen_out). + torch.Tensor: Encoded memory (#batch, maxlen_in, size). + torch.Tensor: Encoded memory mask (#batch, maxlen_in). + + """ + residual = tgt + if self.normalize_before: + tgt = self.norm1(tgt) + tgt = self.feed_forward(tgt) + + x = tgt + if self.self_attn: + if self.normalize_before: + tgt = self.norm2(tgt) + x, fsmn_cache = self.self_attn(tgt, None, fsmn_cache) + x = residual + self.dropout(x) + + if self.src_attn is not None: + residual = x + if self.normalize_before: + x = self.norm3(x) + + x, opt_cache = self.src_attn.forward_chunk(x, memory, opt_cache, chunk_size, look_back) + x = residual + x + + return x, memory, fsmn_cache, opt_cache + + +@tables.register("decoder_classes", "FsmnDecoderSCAMAOpt") +class FsmnDecoderSCAMAOpt(BaseTransformerDecoder): + """ + Author: Shiliang Zhang, Zhifu Gao, Haoneng Luo, Ming Lei, Jie Gao, Zhijie Yan, Lei Xie + SCAMA: Streaming chunk-aware multihead attention for online end-to-end speech recognition + https://arxiv.org/abs/2006.01712 + """ + + def __init__( + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + att_layer_num: int = 6, + kernel_size: int = 21, + sanm_shfit: int = None, + concat_embeds: bool = False, + attention_dim: int = None, + tf2torch_tensor_name_prefix_torch: str = "decoder", + tf2torch_tensor_name_prefix_tf: str = "seq2seq/decoder", + embed_tensor_name_prefix_tf: str = None, + ): + super().__init__( + vocab_size=vocab_size, + encoder_output_size=encoder_output_size, + dropout_rate=dropout_rate, + positional_dropout_rate=positional_dropout_rate, + input_layer=input_layer, + use_output_layer=use_output_layer, + pos_enc_class=pos_enc_class, + normalize_before=normalize_before, + ) + if attention_dim is None: + attention_dim = encoder_output_size + + if input_layer == "embed": + self.embed = torch.nn.Sequential( + torch.nn.Embedding(vocab_size, attention_dim), + ) + elif input_layer == "linear": + self.embed = torch.nn.Sequential( + torch.nn.Linear(vocab_size, attention_dim), + torch.nn.LayerNorm(attention_dim), + torch.nn.Dropout(dropout_rate), + torch.nn.ReLU(), + pos_enc_class(attention_dim, positional_dropout_rate), + ) + else: + raise ValueError(f"only 'embed' or 'linear' is supported: {input_layer}") + + self.normalize_before = normalize_before + if self.normalize_before: + self.after_norm = LayerNorm(attention_dim) + if use_output_layer: + self.output_layer = torch.nn.Linear(attention_dim, vocab_size) + else: + self.output_layer = None + + self.att_layer_num = att_layer_num + self.num_blocks = num_blocks + if sanm_shfit is None: + sanm_shfit = (kernel_size - 1) // 2 + self.decoders = repeat( + att_layer_num, + lambda lnum: DecoderLayerSANM( + attention_dim, + MultiHeadedAttentionSANMDecoder( + attention_dim, self_attention_dropout_rate, kernel_size, sanm_shfit=sanm_shfit + ), + MultiHeadedAttentionCrossAtt( + attention_heads, + attention_dim, + src_attention_dropout_rate, + encoder_output_size=encoder_output_size, + ), + PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate), + dropout_rate, + normalize_before, + concat_after, + ), + ) + if num_blocks - att_layer_num <= 0: + self.decoders2 = None + else: + self.decoders2 = repeat( + num_blocks - att_layer_num, + lambda lnum: DecoderLayerSANM( + attention_dim, + MultiHeadedAttentionSANMDecoder( + attention_dim, + self_attention_dropout_rate, + kernel_size, + sanm_shfit=sanm_shfit, + ), + None, + PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate), + dropout_rate, + normalize_before, + concat_after, + ), + ) + + self.decoders3 = repeat( + 1, + lambda lnum: DecoderLayerSANM( + attention_dim, + None, + None, + PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate), + dropout_rate, + normalize_before, + concat_after, + ), + ) + if concat_embeds: + self.embed_concat_ffn = repeat( + 1, + lambda lnum: DecoderLayerSANM( + attention_dim + encoder_output_size, + None, + None, + PositionwiseFeedForwardDecoderSANM( + attention_dim + encoder_output_size, + linear_units, + dropout_rate, + adim=attention_dim, + ), + dropout_rate, + normalize_before, + concat_after, + ), + ) + else: + self.embed_concat_ffn = None + self.concat_embeds = concat_embeds + self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch + self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf + self.embed_tensor_name_prefix_tf = embed_tensor_name_prefix_tf + + def forward( + self, + hs_pad: torch.Tensor, + hlens: torch.Tensor, + ys_in_pad: torch.Tensor, + ys_in_lens: torch.Tensor, + chunk_mask: torch.Tensor = None, + pre_acoustic_embeds: torch.Tensor = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Forward decoder. + + Args: + hs_pad: encoded memory, float32 (batch, maxlen_in, feat) + hlens: (batch) + ys_in_pad: + input token ids, int64 (batch, maxlen_out) + if input_layer == "embed" + input tensor (batch, maxlen_out, #mels) in the other cases + ys_in_lens: (batch) + Returns: + (tuple): tuple containing: + + x: decoded token score before softmax (batch, maxlen_out, token) + if use_output_layer is True, + olens: (batch, ) + """ + tgt = ys_in_pad + tgt_mask = myutils.sequence_mask(ys_in_lens, device=tgt.device)[:, :, None] + + memory = hs_pad + memory_mask = myutils.sequence_mask(hlens, device=memory.device)[:, None, :] + if chunk_mask is not None: + memory_mask = memory_mask * chunk_mask + if tgt_mask.size(1) != memory_mask.size(1): + memory_mask = torch.cat((memory_mask, memory_mask[:, -2:-1, :]), dim=1) + + x = self.embed(tgt) + + if pre_acoustic_embeds is not None and self.concat_embeds: + x = torch.cat((x, pre_acoustic_embeds), dim=-1) + x, _, _, _, _ = self.embed_concat_ffn(x, None, None, None, None) + + x, tgt_mask, memory, memory_mask, _ = self.decoders(x, tgt_mask, memory, memory_mask) + if self.decoders2 is not None: + x, tgt_mask, memory, memory_mask, _ = self.decoders2(x, tgt_mask, memory, memory_mask) + x, tgt_mask, memory, memory_mask, _ = self.decoders3(x, tgt_mask, memory, memory_mask) + if self.normalize_before: + x = self.after_norm(x) + if self.output_layer is not None: + x = self.output_layer(x) + + olens = tgt_mask.sum(1) + return x, olens + + def score( + self, + ys, + state, + x, + x_mask=None, + pre_acoustic_embeds: torch.Tensor = None, + ): + """Score.""" + ys_mask = myutils.sequence_mask( + torch.tensor([len(ys)], dtype=torch.int32), device=x.device + )[:, :, None] + logp, state = self.forward_one_step( + ys.unsqueeze(0), + ys_mask, + x.unsqueeze(0), + memory_mask=x_mask, + pre_acoustic_embeds=pre_acoustic_embeds, + cache=state, + ) + return logp.squeeze(0), state + + def forward_one_step( + self, + tgt: torch.Tensor, + tgt_mask: torch.Tensor, + memory: torch.Tensor, + memory_mask: torch.Tensor = None, + pre_acoustic_embeds: torch.Tensor = None, + cache: List[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, List[torch.Tensor]]: + """Forward one step. + + Args: + tgt: input token ids, int64 (batch, maxlen_out) + tgt_mask: input token mask, (batch, maxlen_out) + dtype=torch.uint8 in PyTorch 1.2- + dtype=torch.bool in PyTorch 1.2+ (include 1.2) + memory: encoded memory, float32 (batch, maxlen_in, feat) + cache: cached output list of (batch, max_time_out-1, size) + Returns: + y, cache: NN output value and cache per `self.decoders`. + y.shape` is (batch, maxlen_out, token) + """ + + x = tgt[:, -1:] + tgt_mask = None + x = self.embed(x) + + if pre_acoustic_embeds is not None and self.concat_embeds: + x = torch.cat((x, pre_acoustic_embeds), dim=-1) + x, _, _, _, _ = self.embed_concat_ffn(x, None, None, None, None) + + if cache is None: + cache_layer_num = len(self.decoders) + if self.decoders2 is not None: + cache_layer_num += len(self.decoders2) + cache = [None] * cache_layer_num + new_cache = [] + # for c, decoder in zip(cache, self.decoders): + for i in range(self.att_layer_num): + decoder = self.decoders[i] + c = cache[i] + x, tgt_mask, memory, memory_mask, c_ret = decoder.forward_one_step( + x, tgt_mask, memory, memory_mask, cache=c + ) + new_cache.append(c_ret) + + if self.num_blocks - self.att_layer_num >= 1: + for i in range(self.num_blocks - self.att_layer_num): + j = i + self.att_layer_num + decoder = self.decoders2[i] + c = cache[j] + x, tgt_mask, memory, memory_mask, c_ret = decoder.forward_one_step( + x, tgt_mask, memory, memory_mask, cache=c + ) + new_cache.append(c_ret) + + for decoder in self.decoders3: + x, tgt_mask, memory, memory_mask, _ = decoder.forward_one_step( + x, tgt_mask, memory, None, cache=None + ) + + if self.normalize_before: + y = self.after_norm(x[:, -1]) + else: + y = x[:, -1] + if self.output_layer is not None: + y = self.output_layer(y) + y = torch.log_softmax(y, dim=-1) + + return y, new_cache diff --git a/funasr/models/scama/encoder.py b/funasr/models/scama/encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..e1fe9242e056ce6badc0d989d7ff8e66fe8cff5a --- /dev/null +++ b/funasr/models/scama/encoder.py @@ -0,0 +1,499 @@ +#!/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) + +from typing import List +from typing import Optional +from typing import Sequence +from typing import Tuple +from typing import Union +import logging +import torch +import torch.nn as nn +import torch.nn.functional as F +from funasr.models.scama.chunk_utilis import overlap_chunk +import numpy as np +from funasr.train_utils.device_funcs import to_device +from funasr.models.transformer.utils.nets_utils import make_pad_mask +from funasr.models.sanm.attention import MultiHeadedAttention, MultiHeadedAttentionSANM +from funasr.models.transformer.embedding import ( + SinusoidalPositionEncoder, + StreamSinusoidalPositionEncoder, +) +from funasr.models.transformer.layer_norm import LayerNorm +from funasr.models.transformer.utils.multi_layer_conv import Conv1dLinear +from funasr.models.transformer.utils.multi_layer_conv import MultiLayeredConv1d +from funasr.models.transformer.positionwise_feed_forward import ( + PositionwiseFeedForward, # noqa: H301 +) +from funasr.models.transformer.utils.repeat import repeat +from funasr.models.transformer.utils.subsampling import Conv2dSubsampling +from funasr.models.transformer.utils.subsampling import Conv2dSubsampling2 +from funasr.models.transformer.utils.subsampling import Conv2dSubsampling6 +from funasr.models.transformer.utils.subsampling import Conv2dSubsampling8 +from funasr.models.transformer.utils.subsampling import TooShortUttError +from funasr.models.transformer.utils.subsampling import check_short_utt +from funasr.models.transformer.utils.mask import subsequent_mask, vad_mask + +from funasr.models.ctc.ctc import CTC + +from funasr.register import tables + + +class EncoderLayerSANM(nn.Module): + def __init__( + self, + in_size, + size, + self_attn, + feed_forward, + dropout_rate, + normalize_before=True, + concat_after=False, + stochastic_depth_rate=0.0, + ): + """Construct an EncoderLayer object.""" + super(EncoderLayerSANM, self).__init__() + self.self_attn = self_attn + self.feed_forward = feed_forward + self.norm1 = LayerNorm(in_size) + self.norm2 = LayerNorm(size) + self.dropout = nn.Dropout(dropout_rate) + self.in_size = in_size + self.size = size + self.normalize_before = normalize_before + self.concat_after = concat_after + if self.concat_after: + self.concat_linear = nn.Linear(size + size, size) + self.stochastic_depth_rate = stochastic_depth_rate + self.dropout_rate = dropout_rate + + def forward(self, x, mask, cache=None, mask_shfit_chunk=None, mask_att_chunk_encoder=None): + """Compute encoded features. + + Args: + x_input (torch.Tensor): Input tensor (#batch, time, size). + mask (torch.Tensor): Mask tensor for the input (#batch, time). + cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size). + + Returns: + torch.Tensor: Output tensor (#batch, time, size). + torch.Tensor: Mask tensor (#batch, time). + + """ + skip_layer = False + # with stochastic depth, residual connection `x + f(x)` becomes + # `x <- x + 1 / (1 - p) * f(x)` at training time. + stoch_layer_coeff = 1.0 + if self.training and self.stochastic_depth_rate > 0: + skip_layer = torch.rand(1).item() < self.stochastic_depth_rate + stoch_layer_coeff = 1.0 / (1 - self.stochastic_depth_rate) + + if skip_layer: + if cache is not None: + x = torch.cat([cache, x], dim=1) + return x, mask + + residual = x + if self.normalize_before: + x = self.norm1(x) + + if self.concat_after: + x_concat = torch.cat( + ( + x, + self.self_attn( + x, + mask, + mask_shfit_chunk=mask_shfit_chunk, + mask_att_chunk_encoder=mask_att_chunk_encoder, + ), + ), + dim=-1, + ) + if self.in_size == self.size: + x = residual + stoch_layer_coeff * self.concat_linear(x_concat) + else: + x = stoch_layer_coeff * self.concat_linear(x_concat) + else: + if self.in_size == self.size: + x = residual + stoch_layer_coeff * self.dropout( + self.self_attn( + x, + mask, + mask_shfit_chunk=mask_shfit_chunk, + mask_att_chunk_encoder=mask_att_chunk_encoder, + ) + ) + else: + x = stoch_layer_coeff * self.dropout( + self.self_attn( + x, + mask, + mask_shfit_chunk=mask_shfit_chunk, + mask_att_chunk_encoder=mask_att_chunk_encoder, + ) + ) + if not self.normalize_before: + x = self.norm1(x) + + residual = x + if self.normalize_before: + x = self.norm2(x) + x = residual + stoch_layer_coeff * self.dropout(self.feed_forward(x)) + if not self.normalize_before: + x = self.norm2(x) + + return x, mask, cache, mask_shfit_chunk, mask_att_chunk_encoder + + def forward_chunk(self, x, cache=None, chunk_size=None, look_back=0): + """Compute encoded features. + + Args: + x_input (torch.Tensor): Input tensor (#batch, time, size). + mask (torch.Tensor): Mask tensor for the input (#batch, time). + cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size). + + Returns: + torch.Tensor: Output tensor (#batch, time, size). + torch.Tensor: Mask tensor (#batch, time). + + """ + + residual = x + if self.normalize_before: + x = self.norm1(x) + + if self.in_size == self.size: + attn, cache = self.self_attn.forward_chunk(x, cache, chunk_size, look_back) + x = residual + attn + else: + x, cache = self.self_attn.forward_chunk(x, cache, chunk_size, look_back) + + if not self.normalize_before: + x = self.norm1(x) + + residual = x + if self.normalize_before: + x = self.norm2(x) + x = residual + self.feed_forward(x) + if not self.normalize_before: + x = self.norm2(x) + + return x, cache + + +@tables.register("encoder_classes", "SANMEncoderChunkOpt") +class SANMEncoderChunkOpt(nn.Module): + """ + Author: Shiliang Zhang, Zhifu Gao, Haoneng Luo, Ming Lei, Jie Gao, Zhijie Yan, Lei Xie + SCAMA: Streaming chunk-aware multihead attention for online end-to-end speech recognition + https://arxiv.org/abs/2006.01712 + """ + + def __init__( + self, + input_size: int, + output_size: int = 256, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + attention_dropout_rate: float = 0.0, + input_layer: Optional[str] = "conv2d", + pos_enc_class=SinusoidalPositionEncoder, + normalize_before: bool = True, + concat_after: bool = False, + positionwise_layer_type: str = "linear", + positionwise_conv_kernel_size: int = 1, + padding_idx: int = -1, + interctc_layer_idx: List[int] = [], + interctc_use_conditioning: bool = False, + kernel_size: int = 11, + sanm_shfit: int = 0, + selfattention_layer_type: str = "sanm", + chunk_size: Union[int, Sequence[int]] = (16,), + stride: Union[int, Sequence[int]] = (10,), + pad_left: Union[int, Sequence[int]] = (0,), + encoder_att_look_back_factor: Union[int, Sequence[int]] = (1,), + decoder_att_look_back_factor: Union[int, Sequence[int]] = (1,), + tf2torch_tensor_name_prefix_torch: str = "encoder", + tf2torch_tensor_name_prefix_tf: str = "seq2seq/encoder", + ): + super().__init__() + self._output_size = output_size + + if input_layer == "linear": + self.embed = torch.nn.Sequential( + torch.nn.Linear(input_size, output_size), + torch.nn.LayerNorm(output_size), + torch.nn.Dropout(dropout_rate), + torch.nn.ReLU(), + pos_enc_class(output_size, positional_dropout_rate), + ) + elif input_layer == "conv2d": + self.embed = Conv2dSubsampling(input_size, output_size, dropout_rate) + elif input_layer == "conv2d2": + self.embed = Conv2dSubsampling2(input_size, output_size, dropout_rate) + elif input_layer == "conv2d6": + self.embed = Conv2dSubsampling6(input_size, output_size, dropout_rate) + elif input_layer == "conv2d8": + self.embed = Conv2dSubsampling8(input_size, output_size, dropout_rate) + elif input_layer == "embed": + self.embed = torch.nn.Sequential( + torch.nn.Embedding(input_size, output_size, padding_idx=padding_idx), + pos_enc_class(output_size, positional_dropout_rate), + ) + elif input_layer is None: + if input_size == output_size: + self.embed = None + else: + self.embed = torch.nn.Linear(input_size, output_size) + elif input_layer == "pe": + self.embed = SinusoidalPositionEncoder() + elif input_layer == "pe_online": + self.embed = StreamSinusoidalPositionEncoder() + else: + raise ValueError("unknown input_layer: " + input_layer) + self.normalize_before = normalize_before + if positionwise_layer_type == "linear": + positionwise_layer = PositionwiseFeedForward + positionwise_layer_args = ( + output_size, + linear_units, + dropout_rate, + ) + elif positionwise_layer_type == "conv1d": + positionwise_layer = MultiLayeredConv1d + positionwise_layer_args = ( + output_size, + linear_units, + positionwise_conv_kernel_size, + dropout_rate, + ) + elif positionwise_layer_type == "conv1d-linear": + positionwise_layer = Conv1dLinear + positionwise_layer_args = ( + output_size, + linear_units, + positionwise_conv_kernel_size, + dropout_rate, + ) + else: + raise NotImplementedError("Support only linear or conv1d.") + + if selfattention_layer_type == "selfattn": + encoder_selfattn_layer = MultiHeadedAttention + encoder_selfattn_layer_args = ( + attention_heads, + output_size, + attention_dropout_rate, + ) + elif selfattention_layer_type == "sanm": + encoder_selfattn_layer = MultiHeadedAttentionSANM + encoder_selfattn_layer_args0 = ( + attention_heads, + input_size, + output_size, + attention_dropout_rate, + kernel_size, + sanm_shfit, + ) + + encoder_selfattn_layer_args = ( + attention_heads, + output_size, + output_size, + attention_dropout_rate, + kernel_size, + sanm_shfit, + ) + self.encoders0 = repeat( + 1, + lambda lnum: EncoderLayerSANM( + input_size, + output_size, + encoder_selfattn_layer(*encoder_selfattn_layer_args0), + positionwise_layer(*positionwise_layer_args), + dropout_rate, + normalize_before, + concat_after, + ), + ) + + self.encoders = repeat( + num_blocks - 1, + lambda lnum: EncoderLayerSANM( + output_size, + output_size, + encoder_selfattn_layer(*encoder_selfattn_layer_args), + positionwise_layer(*positionwise_layer_args), + dropout_rate, + normalize_before, + concat_after, + ), + ) + if self.normalize_before: + self.after_norm = LayerNorm(output_size) + + self.interctc_layer_idx = interctc_layer_idx + if len(interctc_layer_idx) > 0: + assert 0 < min(interctc_layer_idx) and max(interctc_layer_idx) < num_blocks + self.interctc_use_conditioning = interctc_use_conditioning + self.conditioning_layer = None + shfit_fsmn = (kernel_size - 1) // 2 + self.overlap_chunk_cls = overlap_chunk( + chunk_size=chunk_size, + stride=stride, + pad_left=pad_left, + shfit_fsmn=shfit_fsmn, + encoder_att_look_back_factor=encoder_att_look_back_factor, + decoder_att_look_back_factor=decoder_att_look_back_factor, + ) + self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch + self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf + + def output_size(self) -> int: + return self._output_size + + def forward( + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + prev_states: torch.Tensor = None, + ctc: CTC = None, + ind: int = 0, + ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Embed positions in tensor. + + Args: + xs_pad: input tensor (B, L, D) + ilens: input length (B) + prev_states: Not to be used now. + Returns: + position embedded tensor and mask + """ + masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device) + xs_pad *= self.output_size() ** 0.5 + if self.embed is None: + xs_pad = xs_pad + elif ( + isinstance(self.embed, Conv2dSubsampling) + or isinstance(self.embed, Conv2dSubsampling2) + or isinstance(self.embed, Conv2dSubsampling6) + or isinstance(self.embed, Conv2dSubsampling8) + ): + short_status, limit_size = check_short_utt(self.embed, xs_pad.size(1)) + if short_status: + raise TooShortUttError( + f"has {xs_pad.size(1)} frames and is too short for subsampling " + + f"(it needs more than {limit_size} frames), return empty results", + xs_pad.size(1), + limit_size, + ) + xs_pad, masks = self.embed(xs_pad, masks) + else: + xs_pad = self.embed(xs_pad) + + mask_shfit_chunk, mask_att_chunk_encoder = None, None + if self.overlap_chunk_cls is not None: + ilens = masks.squeeze(1).sum(1) + chunk_outs = self.overlap_chunk_cls.gen_chunk_mask(ilens, ind) + xs_pad, ilens = self.overlap_chunk_cls.split_chunk(xs_pad, ilens, chunk_outs=chunk_outs) + masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device) + mask_shfit_chunk = self.overlap_chunk_cls.get_mask_shfit_chunk( + chunk_outs, xs_pad.device, xs_pad.size(0), dtype=xs_pad.dtype + ) + mask_att_chunk_encoder = self.overlap_chunk_cls.get_mask_att_chunk_encoder( + chunk_outs, xs_pad.device, xs_pad.size(0), dtype=xs_pad.dtype + ) + + encoder_outs = self.encoders0(xs_pad, masks, None, mask_shfit_chunk, mask_att_chunk_encoder) + xs_pad, masks = encoder_outs[0], encoder_outs[1] + intermediate_outs = [] + if len(self.interctc_layer_idx) == 0: + encoder_outs = self.encoders( + xs_pad, masks, None, mask_shfit_chunk, mask_att_chunk_encoder + ) + xs_pad, masks = encoder_outs[0], encoder_outs[1] + else: + for layer_idx, encoder_layer in enumerate(self.encoders): + encoder_outs = encoder_layer( + xs_pad, masks, None, mask_shfit_chunk, mask_att_chunk_encoder + ) + xs_pad, masks = encoder_outs[0], encoder_outs[1] + if layer_idx + 1 in self.interctc_layer_idx: + encoder_out = xs_pad + + # intermediate outputs are also normalized + if self.normalize_before: + encoder_out = self.after_norm(encoder_out) + + intermediate_outs.append((layer_idx + 1, encoder_out)) + + if self.interctc_use_conditioning: + ctc_out = ctc.softmax(encoder_out) + xs_pad = xs_pad + self.conditioning_layer(ctc_out) + + if self.normalize_before: + xs_pad = self.after_norm(xs_pad) + + olens = masks.squeeze(1).sum(1) + if len(intermediate_outs) > 0: + return (xs_pad, intermediate_outs), olens, None + return xs_pad, olens, None + + def _add_overlap_chunk(self, feats: np.ndarray, cache: dict = {}): + if len(cache) == 0: + return feats + cache["feats"] = to_device(cache["feats"], device=feats.device) + overlap_feats = torch.cat((cache["feats"], feats), dim=1) + cache["feats"] = overlap_feats[:, -(cache["chunk_size"][0] + cache["chunk_size"][2]) :, :] + return overlap_feats + + def forward_chunk( + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + cache: dict = None, + **kwargs, + ): + is_final = kwargs.get("is_final", False) + xs_pad *= self.output_size() ** 0.5 + if self.embed is None: + xs_pad = xs_pad + else: + xs_pad = self.embed(xs_pad, cache) + if cache["tail_chunk"]: + xs_pad = to_device(cache["feats"], device=xs_pad.device) + else: + xs_pad = self._add_overlap_chunk(xs_pad, cache) + if cache["opt"] is None: + cache_layer_num = len(self.encoders0) + len(self.encoders) + new_cache = [None] * cache_layer_num + else: + new_cache = cache["opt"] + + for layer_idx, encoder_layer in enumerate(self.encoders0): + encoder_outs = encoder_layer.forward_chunk( + xs_pad, new_cache[layer_idx], cache["chunk_size"], cache["encoder_chunk_look_back"] + ) + xs_pad, new_cache[0] = encoder_outs[0], encoder_outs[1] + + for layer_idx, encoder_layer in enumerate(self.encoders): + encoder_outs = encoder_layer.forward_chunk( + xs_pad, + new_cache[layer_idx + len(self.encoders0)], + cache["chunk_size"], + cache["encoder_chunk_look_back"], + ) + xs_pad, new_cache[layer_idx + len(self.encoders0)] = encoder_outs[0], encoder_outs[1] + + if self.normalize_before: + xs_pad = self.after_norm(xs_pad) + if cache["encoder_chunk_look_back"] > 0 or cache["encoder_chunk_look_back"] == -1: + cache["opt"] = new_cache + + return xs_pad, ilens, None diff --git a/funasr/models/scama/model.py b/funasr/models/scama/model.py new file mode 100644 index 0000000000000000000000000000000000000000..c15f435a46694cc30399baaf811dd470aec8c6c9 --- /dev/null +++ b/funasr/models/scama/model.py @@ -0,0 +1,739 @@ +#!/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 torch.nn as nn +import torch.functional as F +import logging +from typing import Dict, Tuple +from contextlib import contextmanager +from distutils.version import LooseVersion + +from funasr.register import tables +from funasr.models.ctc.ctc import CTC +from funasr.utils import postprocess_utils +from funasr.metrics.compute_acc import th_accuracy +from funasr.utils.datadir_writer import DatadirWriter +from funasr.models.paraformer.model import Paraformer +from funasr.models.paraformer.search import Hypothesis +from funasr.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 + +if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): + from torch.cuda.amp import autocast +else: + # Nothing to do if torch<1.6.0 + @contextmanager + def autocast(enabled=True): + yield + + +@tables.register("model_classes", "SCAMA") +class SCAMA(nn.Module): + """ + Author: Shiliang Zhang, Zhifu Gao, Haoneng Luo, Ming Lei, Jie Gao, Zhijie Yan, Lei Xie + SCAMA: Streaming chunk-aware multihead attention for online end-to-end speech recognition + https://arxiv.org/abs/2006.01712 + """ + + def __init__( + self, + specaug: str = None, + specaug_conf: dict = None, + normalize: str = None, + normalize_conf: dict = None, + encoder: str = None, + encoder_conf: dict = None, + decoder: str = None, + decoder_conf: dict = None, + ctc: str = None, + ctc_conf: dict = None, + ctc_weight: float = 0.5, + predictor: str = None, + predictor_conf: dict = None, + predictor_bias: int = 0, + predictor_weight: float = 0.0, + 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, + ) + if ctc_weight > 0.0: + + if ctc_conf is None: + ctc_conf = {} + + ctc = CTC(odim=vocab_size, encoder_output_size=encoder_output_size, **ctc_conf) + + predictor_class = tables.predictor_classes.get(predictor) + predictor = predictor_class(**predictor_conf) + + # note that eos is the same as sos (equivalent ID) + self.blank_id = blank_id + self.sos = sos if sos is not None else vocab_size - 1 + self.eos = eos if eos is not None else vocab_size - 1 + self.vocab_size = vocab_size + self.ignore_id = ignore_id + self.ctc_weight = ctc_weight + + self.specaug = specaug + self.normalize = normalize + + self.encoder = encoder + + if ctc_weight == 1.0: + self.decoder = None + else: + self.decoder = decoder + + self.criterion_att = LabelSmoothingLoss( + size=vocab_size, + padding_idx=ignore_id, + smoothing=lsm_weight, + normalize_length=length_normalized_loss, + ) + + if ctc_weight == 0.0: + self.ctc = None + else: + self.ctc = ctc + + self.predictor = predictor + self.predictor_weight = predictor_weight + self.predictor_bias = predictor_bias + + self.criterion_pre = mae_loss(normalize_length=length_normalized_loss) + + self.share_embedding = share_embedding + if self.share_embedding: + self.decoder.embed = None + + self.length_normalized_loss = length_normalized_loss + self.beam_search = None + self.error_calculator = None + + 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 = kwargs.get("decoder_attention_chunk_type", "chunk") + + def forward( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + text: torch.Tensor, + text_lengths: torch.Tensor, + **kwargs, + ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]: + """Encoder + Decoder + Calc loss + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + text: (Batch, Length) + text_lengths: (Batch,) + """ + + decoding_ind = kwargs.get("decoding_ind") + if len(text_lengths.size()) > 1: + text_lengths = text_lengths[:, 0] + if len(speech_lengths.size()) > 1: + speech_lengths = speech_lengths[:, 0] + + batch_size = speech.shape[0] + + # Encoder + ind = self.encoder.overlap_chunk_cls.random_choice(self.training, decoding_ind) + encoder_out, encoder_out_lens = self.encode(speech, speech_lengths, ind=ind) + + loss_ctc, cer_ctc = None, None + loss_pre = None + stats = dict() + + # decoder: CTC branch + + if self.ctc_weight > 0.0: + + encoder_out_ctc, encoder_out_lens_ctc = self.encoder.overlap_chunk_cls.remove_chunk( + encoder_out, encoder_out_lens, chunk_outs=None + ) + + loss_ctc, cer_ctc = self._calc_ctc_loss( + encoder_out_ctc, encoder_out_lens_ctc, text, text_lengths + ) + # Collect CTC branch stats + stats["loss_ctc"] = loss_ctc.detach() if loss_ctc is not None else None + stats["cer_ctc"] = cer_ctc + + # decoder: Attention decoder branch + loss_att, acc_att, cer_att, wer_att, loss_pre = self._calc_att_predictor_loss( + encoder_out, encoder_out_lens, text, text_lengths + ) + + # 3. CTC-Att loss definition + if self.ctc_weight == 0.0: + loss = loss_att + loss_pre * self.predictor_weight + else: + loss = ( + self.ctc_weight * loss_ctc + + (1 - self.ctc_weight) * 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 + + stats["loss"] = torch.clone(loss.detach()) + + # force_gatherable: to-device and to-tensor if scalar for DataParallel + if self.length_normalized_loss: + batch_size = (text_lengths + self.predictor_bias).sum() + loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) + return loss, stats, weight + + def encode( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Encoder. Note that this method is used by asr_inference.py + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + ind: int + """ + with autocast(False): + + # Data augmentation + if self.specaug is not None and self.training: + speech, speech_lengths = self.specaug(speech, speech_lengths) + + # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN + if self.normalize is not None: + speech, speech_lengths = self.normalize(speech, speech_lengths) + + # Forward encoder + encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths) + if isinstance(encoder_out, tuple): + encoder_out = encoder_out[0] + + return encoder_out, encoder_out_lens + + def encode_chunk( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + cache: dict = None, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Frontend + Encoder. Note that this method is used by asr_inference.py + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + ind: int + """ + with autocast(False): + + # Data augmentation + if self.specaug is not None and self.training: + speech, speech_lengths = self.specaug(speech, speech_lengths) + + # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN + if self.normalize is not None: + speech, speech_lengths = self.normalize(speech, speech_lengths) + + # Forward encoder + encoder_out, encoder_out_lens, _ = self.encoder.forward_chunk( + speech, speech_lengths, cache=cache["encoder"] + ) + if isinstance(encoder_out, tuple): + encoder_out = encoder_out[0] + + return encoder_out, torch.tensor([encoder_out.size(1)]) + + def calc_predictor_chunk(self, encoder_out, encoder_out_lens, cache=None, **kwargs): + is_final = kwargs.get("is_final", False) + + return self.predictor.forward_chunk(encoder_out, cache["encoder"], is_final=is_final) + + 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 + ) + + 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, + ) + + # 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_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 + + 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 + ) + + 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, + ) + + return ( + pre_acoustic_embeds, + pre_token_length, + predictor_alignments, + predictor_alignments_len, + scama_mask, + ) + + def init_beam_search( + self, + **kwargs, + ): + + from funasr.models.scama.beam_search import BeamSearchScamaStreaming + + from funasr.models.transformer.scorers.ctc import CTCPrefixScorer + from funasr.models.transformer.scorers.length_bonus import LengthBonus + + # 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=self.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 = BeamSearchScamaStreaming( + beam_size=kwargs.get("beam_size", 2), + 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", + ) + # beam_search.to(device=kwargs.get("device", "cpu"), dtype=getattr(torch, kwargs.get("dtype", "float32"))).eval() + # for scorer in scorers.values(): + # if isinstance(scorer, torch.nn.Module): + # scorer.to(device=kwargs.get("device", "cpu"), dtype=getattr(torch, kwargs.get("dtype", "float32"))).eval() + self.beam_search = beam_search + + def generate_chunk( + self, + speech, + speech_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): + cache = kwargs.get("cache", {}) + speech = speech.to(device=kwargs["device"]) + speech_lengths = speech_lengths.to(device=kwargs["device"]) + + # Encoder + encoder_out, encoder_out_lens = self.encode_chunk( + speech, speech_lengths, cache=cache, is_final=kwargs.get("is_final", False) + ) + if isinstance(encoder_out, tuple): + encoder_out = encoder_out[0] + if "running_hyps" not in cache: + running_hyps = self.beam_search.init_hyp(encoder_out) + cache["running_hyps"] = running_hyps + + # predictor + predictor_outs = self.calc_predictor_chunk( + encoder_out, + encoder_out_lens, + cache=cache, + is_final=kwargs.get("is_final", False), + ) + pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = ( + predictor_outs[0], + predictor_outs[1], + predictor_outs[2], + predictor_outs[3], + ) + pre_token_length = pre_token_length.round().long() + + if torch.max(pre_token_length) < 1: + return [] + maxlen = minlen = pre_token_length + if kwargs.get("is_final", False): + maxlen += kwargs.get("token_num_relax", 5) + minlen = max(0, minlen - kwargs.get("token_num_relax", 5)) + # c. Passed the encoder result and the beam search + nbest_hyps = self.beam_search( + x=encoder_out[0], + scama_mask=None, + pre_acoustic_embeds=pre_acoustic_embeds, + maxlen=int(maxlen), + minlen=int(minlen), + cache=cache, + ) + + cache["running_hyps"] = nbest_hyps + nbest_hyps = nbest_hyps[: self.nbest] + + results = [] + for hyp in nbest_hyps: + # assert isinstance(hyp, (Hypothesis)), type(hyp) + + # remove sos/eos and get results + last_pos = -1 + if isinstance(hyp.yseq, list): + token_int = hyp.yseq[1:last_pos] + else: + token_int = hyp.yseq[1:last_pos].tolist() + + # remove blank symbol id, which is assumed to be 0 + token_int = list( + filter( + lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int + ) + ) + + # Change integer-ids to tokens + token = tokenizer.ids2tokens(token_int) + # text = tokenizer.tokens2text(token) + + result_i = token + + results.extend(result_i) + + return results + + def init_cache(self, cache: dict = {}, **kwargs): + device = kwargs.get("device", "cuda") + + chunk_size = kwargs.get("chunk_size", [0, 10, 5]) + encoder_chunk_look_back = kwargs.get("encoder_chunk_look_back", 0) + decoder_chunk_look_back = kwargs.get("decoder_chunk_look_back", 0) + batch_size = 1 + + enc_output_size = kwargs["encoder_conf"]["output_size"] + feats_dims = kwargs["frontend_conf"]["n_mels"] * kwargs["frontend_conf"]["lfr_m"] + + cache_encoder = { + "start_idx": 0, + "cif_hidden": torch.zeros((batch_size, 1, enc_output_size)).to(device=device), + "cif_alphas": torch.zeros((batch_size, 1)).to(device=device), + "chunk_size": chunk_size, + "encoder_chunk_look_back": encoder_chunk_look_back, + "last_chunk": False, + "opt": None, + "feats": torch.zeros((batch_size, chunk_size[0] + chunk_size[2], feats_dims)).to( + device=device + ), + "tail_chunk": False, + } + cache["encoder"] = cache_encoder + + cache_decoder = { + "decode_fsmn": None, + "decoder_chunk_look_back": decoder_chunk_look_back, + "opt": None, + "chunk_size": chunk_size, + } + cache["decoder"] = cache_decoder + cache["frontend"] = {} + + cache["prev_samples"] = torch.empty(0).to(device=device) + + return cache + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + cache: dict = {}, + **kwargs, + ): + + # init beamsearch + is_use_ctc = kwargs.get("decoding_ctc_weight", 0.0) > 0.00001 and self.ctc != None + is_use_lm = ( + kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None + ) + + if self.beam_search is None: + + logging.info("enable beam_search") + self.init_beam_search(**kwargs) + self.nbest = kwargs.get("nbest", 1) + + if len(cache) == 0: + self.init_cache(cache, **kwargs) + + meta_data = {} + chunk_size = kwargs.get("chunk_size", [0, 10, 5]) + chunk_stride_samples = int(chunk_size[1] * 960) # 600ms + + time1 = time.perf_counter() + cfg = {"is_final": kwargs.get("is_final", False)} + 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, + cache=cfg, + ) + _is_final = cfg["is_final"] # if data_in is a file or url, set is_final=True + + time2 = time.perf_counter() + meta_data["load_data"] = f"{time2 - time1:0.3f}" + assert len(audio_sample_list) == 1, "batch_size must be set 1" + + audio_sample = torch.cat((cache["prev_samples"], audio_sample_list[0])) + + n = int(len(audio_sample) // chunk_stride_samples + int(_is_final)) + m = int(len(audio_sample) % chunk_stride_samples * (1 - int(_is_final))) + tokens = [] + for i in range(n): + kwargs["is_final"] = _is_final and i == n - 1 + audio_sample_i = audio_sample[i * chunk_stride_samples : (i + 1) * chunk_stride_samples] + + # extract fbank feats + speech, speech_lengths = extract_fbank( + [audio_sample_i], + data_type=kwargs.get("data_type", "sound"), + frontend=frontend, + cache=cache["frontend"], + is_final=kwargs["is_final"], + ) + 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 + ) + + tokens_i = self.generate_chunk( + speech, + speech_lengths, + key=key, + tokenizer=tokenizer, + cache=cache, + frontend=frontend, + **kwargs, + ) + tokens.extend(tokens_i) + + text_postprocessed, _ = postprocess_utils.sentence_postprocess(tokens) + + result_i = {"key": key[0], "text": text_postprocessed} + result = [result_i] + + cache["prev_samples"] = audio_sample[:-m] + if _is_final: + self.init_cache(cache, **kwargs) + + if kwargs.get("output_dir"): + writer = DatadirWriter(kwargs.get("output_dir")) + ibest_writer = writer[f"{1}best_recog"] + ibest_writer["token"][key[0]] = " ".join(tokens) + ibest_writer["text"][key[0]] = text_postprocessed + + return result, meta_data diff --git a/funasr/models/scama/template.yaml b/funasr/models/scama/template.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bc2e210b2d3652bd32fd12a198fb623f37616185 --- /dev/null +++ b/funasr/models/scama/template.yaml @@ -0,0 +1,127 @@ +# This is an example that demonstrates how to configure a model file. +# You can modify the configuration according to your own requirements. + +# to print the register_table: +# from funasr.register import tables +# tables.print() + +# network architecture +model: SCAMA +model_conf: + ctc_weight: 0.0 + lsm_weight: 0.1 + length_normalized_loss: true + +# encoder +encoder: SANMEncoderChunkOpt +encoder_conf: + output_size: 512 + attention_heads: 4 + linear_units: 2048 + num_blocks: 50 + dropout_rate: 0.1 + positional_dropout_rate: 0.1 + attention_dropout_rate: 0.1 + input_layer: pe + pos_enc_class: SinusoidalPositionEncoder + normalize_before: true + kernel_size: 11 + sanm_shfit: 0 + selfattention_layer_type: sanm + +# decoder +decoder: FsmnDecoderSCAMAOpt +decoder_conf: + attention_heads: 4 + linear_units: 2048 + num_blocks: 16 + dropout_rate: 0.1 + positional_dropout_rate: 0.1 + self_attention_dropout_rate: 0.1 + src_attention_dropout_rate: 0.1 + att_layer_num: 16 + kernel_size: 11 + sanm_shfit: 0 + +predictor: CifPredictorV2 +predictor_conf: + idim: 512 + threshold: 1.0 + l_order: 1 + r_order: 1 + tail_threshold: 0.45 + +# frontend related +frontend: WavFrontend +frontend_conf: + fs: 16000 + window: hamming + n_mels: 80 + frame_length: 25 + frame_shift: 10 + lfr_m: 7 + lfr_n: 6 + +specaug: SpecAugLFR +specaug_conf: + apply_time_warp: false + time_warp_window: 5 + time_warp_mode: bicubic + apply_freq_mask: true + freq_mask_width_range: + - 0 + - 30 + lfr_rate: 6 + num_freq_mask: 1 + apply_time_mask: true + time_mask_width_range: + - 0 + - 12 + num_time_mask: 1 + +train_conf: + accum_grad: 1 + grad_clip: 5 + max_epoch: 150 + val_scheduler_criterion: + - valid + - acc + best_model_criterion: + - - valid + - acc + - max + keep_nbest_models: 10 + avg_nbest_model: 10 + log_interval: 50 + +optim: adam +optim_conf: + lr: 0.0005 +scheduler: warmuplr +scheduler_conf: + warmup_steps: 30000 + +dataset: AudioDataset +dataset_conf: + index_ds: IndexDSJsonl + batch_sampler: BatchSampler + batch_type: example # example or length + batch_size: 1 # if batch_type is example, batch_size is the numbers of samples; if length, batch_size is source_token_len+target_token_len; + max_token_length: 2048 # filter samples if source_token_len+target_token_len > max_token_length, + buffer_size: 500 + shuffle: True + num_workers: 0 + +tokenizer: CharTokenizer +tokenizer_conf: + unk_symbol: + split_with_space: true + + +ctc_conf: + dropout_rate: 0.0 + ctc_type: builtin + reduce: true + ignore_nan_grad: true + +normalize: null diff --git a/funasr/models/scama/utils.py b/funasr/models/scama/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c3f7bc37c77d2d83567b8a6a30dca631ad1a3aff --- /dev/null +++ b/funasr/models/scama/utils.py @@ -0,0 +1,89 @@ +import os +import yaml +import torch +import numpy as np +from torch.nn import functional as F + + +def sequence_mask(lengths, maxlen=None, dtype=torch.float32, device=None): + if maxlen is None: + maxlen = lengths.max() + row_vector = torch.arange(0, maxlen, 1).to(lengths.device) + matrix = torch.unsqueeze(lengths, dim=-1) + mask = row_vector < matrix + mask = mask.detach() + + return mask.type(dtype).to(device) if device is not None else mask.type(dtype) + + +def apply_cmvn(inputs, mvn): + device = inputs.device + dtype = inputs.dtype + frame, dim = inputs.shape + meams = np.tile(mvn[0:1, :dim], (frame, 1)) + vars = np.tile(mvn[1:2, :dim], (frame, 1)) + inputs -= torch.from_numpy(meams).type(dtype).to(device) + inputs *= torch.from_numpy(vars).type(dtype).to(device) + + return inputs.type(torch.float32) + + +def drop_and_add( + inputs: torch.Tensor, + outputs: torch.Tensor, + training: bool, + dropout_rate: float = 0.1, + stoch_layer_coeff: float = 1.0, +): + + outputs = F.dropout(outputs, p=dropout_rate, training=training, inplace=True) + outputs *= stoch_layer_coeff + + input_dim = inputs.size(-1) + output_dim = outputs.size(-1) + + if input_dim == output_dim: + outputs += inputs + return outputs + + +def proc_tf_vocab(vocab_path): + with open(vocab_path, encoding="utf-8") as f: + token_list = [line.rstrip() for line in f] + if "" not in token_list: + token_list.append("") + return token_list + + +def gen_config_for_tfmodel(config_path, vocab_path, output_dir): + token_list = proc_tf_vocab(vocab_path) + with open(config_path, encoding="utf-8") as f: + config = yaml.safe_load(f) + + config["token_list"] = token_list + + if not os.path.exists(output_dir): + os.makedirs(output_dir) + + with open(os.path.join(output_dir, "config.yaml"), "w", encoding="utf-8") as f: + yaml_no_alias_safe_dump(config, f, indent=4, sort_keys=False) + + +class NoAliasSafeDumper(yaml.SafeDumper): + # Disable anchor/alias in yaml because looks ugly + def ignore_aliases(self, data): + return True + + +def yaml_no_alias_safe_dump(data, stream=None, **kwargs): + """Safe-dump in yaml with no anchor/alias""" + return yaml.dump(data, stream, allow_unicode=True, Dumper=NoAliasSafeDumper, **kwargs) + + +if __name__ == "__main__": + import sys + + config_path = sys.argv[1] + vocab_path = sys.argv[2] + output_dir = sys.argv[3] + gen_config_for_tfmodel(config_path, vocab_path, output_dir) diff --git a/funasr/models/seaco_paraformer/__init__.py b/funasr/models/seaco_paraformer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/models/seaco_paraformer/export_meta.py b/funasr/models/seaco_paraformer/export_meta.py new file mode 100644 index 0000000000000000000000000000000000000000..947f4d9569192605cfb9f60b19a0262ff22e064e --- /dev/null +++ b/funasr/models/seaco_paraformer/export_meta.py @@ -0,0 +1,202 @@ +#!/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 + + +class ContextualEmbedderExport(torch.nn.Module): + def __init__( + self, + model, + max_seq_len=512, + feats_dim=560, + **kwargs, + ): + super().__init__() + self.embedding = model.decoder.embed # model.bias_embed + model.bias_encoder.batch_first = False + self.bias_encoder = model.bias_encoder + + def forward(self, hotword): + hotword = self.embedding(hotword).transpose(0, 1) # batch second + hw_embed, (_, _) = self.bias_encoder(hotword) + return hw_embed + + def export_dummy_inputs(self): + hotword = torch.tensor( + [ + [10, 11, 12, 13, 14, 10, 11, 12, 13, 14], + [100, 101, 0, 0, 0, 0, 0, 0, 0, 0], + [1, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [10, 11, 12, 13, 14, 10, 11, 12, 13, 14], + [100, 101, 0, 0, 0, 0, 0, 0, 0, 0], + [1, 0, 0, 0, 0, 0, 0, 0, 0, 0], + ], + dtype=torch.int32, + ) + # hotword_length = torch.tensor([10, 2, 1], dtype=torch.int32) + return hotword + + def export_input_names(self): + return ["hotword"] + + def export_output_names(self): + return ["hw_embed"] + + def export_dynamic_axes(self): + return { + "hotword": { + 0: "num_hotwords", + }, + "hw_embed": { + 1: "num_hotwords", + }, + } + + def export_name(self): + return "model_eb.onnx" + + +def export_rebuild_model(model, **kwargs): + model.device = kwargs.get("device") + is_onnx = kwargs.get("type", "onnx") == "onnx" + encoder_class = tables.encoder_classes.get(kwargs["encoder"] + "Export") + model.encoder = encoder_class(model.encoder, onnx=is_onnx) + + predictor_class = tables.predictor_classes.get(kwargs["predictor"] + "Export") + model.predictor = predictor_class(model.predictor, onnx=is_onnx) + + # before decoder convert into export class + embedder_class = ContextualEmbedderExport + embedder_model = embedder_class(model, onnx=is_onnx) + + decoder_class = tables.decoder_classes.get(kwargs["decoder"] + "Export") + model.decoder = decoder_class(model.decoder, onnx=is_onnx) + + seaco_decoder_class = tables.decoder_classes.get(kwargs["seaco_decoder"] + "Export") + model.seaco_decoder = seaco_decoder_class(model.seaco_decoder, onnx=is_onnx) + + from funasr.utils.torch_function import sequence_mask + + model.make_pad_mask = sequence_mask(kwargs["max_seq_len"], flip=False) + + from funasr.utils.torch_function import sequence_mask + + model.make_pad_mask = sequence_mask(kwargs["max_seq_len"], flip=False) + model.feats_dim = 560 + model.NOBIAS = 8377 + + import copy + import types + + backbone_model = copy.copy(model) + + # backbone + backbone_model.forward = types.MethodType(export_backbone_forward, backbone_model) + backbone_model.export_dummy_inputs = types.MethodType( + export_backbone_dummy_inputs, backbone_model + ) + backbone_model.export_input_names = types.MethodType( + export_backbone_input_names, backbone_model + ) + backbone_model.export_output_names = types.MethodType( + export_backbone_output_names, backbone_model + ) + backbone_model.export_dynamic_axes = types.MethodType( + export_backbone_dynamic_axes, backbone_model + ) + + embedder_model.export_name = "model_eb" + backbone_model.export_name = "model" + + return backbone_model, embedder_model + + +def export_backbone_forward( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + bias_embed: torch.Tensor, + # lmbd: float, +): + # a. To device + batch = {"speech": speech, "speech_lengths": speech_lengths} + + enc, enc_len = self.encoder(**batch) + mask = self.make_pad_mask(enc_len)[:, None, :] + pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = self.predictor(enc, mask) + pre_token_length = pre_token_length.floor().type(torch.int32) + + decoder_out, decoder_hidden, _ = self.decoder( + enc, enc_len, pre_acoustic_embeds, pre_token_length, return_hidden=True, return_both=True + ) + decoder_out = torch.log_softmax(decoder_out, dim=-1) + # seaco forward + B, N, D = bias_embed.shape + _contextual_length = torch.ones(B) * N + + # ASF + hotword_scores = self.seaco_decoder.forward_asf6( + bias_embed, _contextual_length, decoder_hidden, pre_token_length + ) + hotword_scores = hotword_scores[0].sum(0).sum(0) + # _ = self.decoder2(bias_embed, _contextual_length, decoder_hidden, pre_token_length) + # hotword_scores = self.decoder2.model.decoders[-1].attn_mat[0][0].sum(0).sum(0) + dec_filter = torch.sort(hotword_scores, descending=True)[1][:51] + contextual_info = bias_embed[:, dec_filter] + num_hot_word = contextual_info.shape[1] + _contextual_length = torch.Tensor([num_hot_word]).int().repeat(B).to(enc.device) + + # again + cif_attended, _ = self.seaco_decoder( + contextual_info, _contextual_length, pre_acoustic_embeds, pre_token_length + ) + dec_attended, _ = self.seaco_decoder( + contextual_info, _contextual_length, decoder_hidden, pre_token_length + ) + merged = cif_attended + dec_attended + dha_output = self.hotword_output_layer(merged) + dha_pred = torch.log_softmax(dha_output, dim=-1) + # merging logits + dha_ids = dha_pred.max(-1)[-1] + dha_mask = (dha_ids == self.NOBIAS).int().unsqueeze(-1) + decoder_out = decoder_out * dha_mask + dha_pred * (1 - dha_mask) + + # get predicted timestamps + us_alphas, us_cif_peak = self.predictor.get_upsample_timestmap(enc, mask, pre_token_length) + + return decoder_out, pre_token_length, us_alphas, us_cif_peak + + +def export_backbone_dummy_inputs(self): + speech = torch.randn(2, 30, self.feats_dim) + speech_lengths = torch.tensor([15, 30], dtype=torch.int32) + bias_embed = torch.randn(2, 1, 512) + return (speech, speech_lengths, bias_embed) + + +def export_backbone_input_names(self): + return ["speech", "speech_lengths", "bias_embed"] + + +def export_backbone_output_names(self): + return ["logits", "token_num", "us_alphas", "us_cif_peak"] + + +def export_backbone_dynamic_axes(self): + return { + "speech": {0: "batch_size", 1: "feats_length"}, + "speech_lengths": { + 0: "batch_size", + }, + "bias_embed": {0: "batch_size", 1: "num_hotwords"}, + "logits": {0: "batch_size", 1: "logits_length"}, + "pre_acoustic_embeds": {1: "feats_length1"}, + "us_alphas": {0: "batch_size", 1: "alphas_length"}, + "us_cif_peak": {0: "batch_size", 1: "alphas_length"}, + } + diff --git a/funasr/models/seaco_paraformer/model.py b/funasr/models/seaco_paraformer/model.py new file mode 100644 index 0000000000000000000000000000000000000000..8b8ade866ecd54482408787dadde322c66fc07ec --- /dev/null +++ b/funasr/models/seaco_paraformer/model.py @@ -0,0 +1,661 @@ +#!/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 os +import re +import time +import copy +import torch +import codecs +import logging +import tempfile +import requests +import numpy as np +from typing import Dict, Tuple +from contextlib import contextmanager +from distutils.version import LooseVersion + +from funasr.register import tables +from funasr.utils import postprocess_utils +from funasr.models.paraformer.model import Paraformer +from funasr.utils.datadir_writer import DatadirWriter +from funasr.models.paraformer.search import Hypothesis +from funasr.train_utils.device_funcs import force_gatherable +from funasr.models.bicif_paraformer.model import BiCifParaformer +from funasr.losses.label_smoothing_loss import LabelSmoothingLoss +from funasr.models.transformer.utils.add_sos_eos import add_sos_eos +from funasr.utils.timestamp_tools import ts_prediction_lfr6_standard +from funasr.models.transformer.utils.nets_utils import make_pad_mask, pad_list +from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank + +from funasr.utils.run_bmodel import EngineOV +from funasr.models.bicif_paraformer.cif_predictor import cif +from funasr.models.bicif_paraformer.cif_predictor import cif_wo_hidden + +if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): + from torch.cuda.amp import autocast +else: + # Nothing to do if torch<1.6.0 + @contextmanager + def autocast(enabled=True): + yield + + +@tables.register("model_classes", "SeacoParaformer") +class SeacoParaformer(BiCifParaformer, Paraformer): + """ + Author: Speech Lab of DAMO Academy, Alibaba Group + SeACo-Paraformer: A Non-Autoregressive ASR System with Flexible and Effective Hotword Customization Ability + https://arxiv.org/abs/2308.03266 + """ + + def __init__( + self, + *args, + **kwargs, + ): + super().__init__(*args, **kwargs) + + self.inner_dim = kwargs.get("inner_dim", 256) + self.bias_encoder_type = kwargs.get("bias_encoder_type", "lstm") + bias_encoder_dropout_rate = kwargs.get("bias_encoder_dropout_rate", 0.0) + bias_encoder_bid = kwargs.get("bias_encoder_bid", False) + seaco_lsm_weight = kwargs.get("seaco_lsm_weight", 0.0) + seaco_length_normalized_loss = kwargs.get("seaco_length_normalized_loss", True) + + # bias encoder + if self.bias_encoder_type == "lstm": + self.bias_encoder = torch.nn.LSTM( + self.inner_dim, + self.inner_dim, + 2, + batch_first=True, + dropout=bias_encoder_dropout_rate, + bidirectional=bias_encoder_bid, + ) + if bias_encoder_bid: + self.lstm_proj = torch.nn.Linear(self.inner_dim * 2, self.inner_dim) + else: + self.lstm_proj = None + # self.bias_embed = torch.nn.Embedding(self.vocab_size, self.inner_dim) + elif self.bias_encoder_type == "mean": + self.bias_embed = torch.nn.Embedding(self.vocab_size, self.inner_dim) + else: + logging.error("Unsupport bias encoder type: {}".format(self.bias_encoder_type)) + + # seaco decoder + seaco_decoder = kwargs.get("seaco_decoder", None) + if seaco_decoder is not None: + seaco_decoder_conf = kwargs.get("seaco_decoder_conf") + seaco_decoder_class = tables.decoder_classes.get(seaco_decoder) + self.seaco_decoder = seaco_decoder_class( + vocab_size=self.vocab_size, + encoder_output_size=self.inner_dim, + **seaco_decoder_conf, + ) + self.hotword_output_layer = torch.nn.Linear(self.inner_dim, self.vocab_size) + self.criterion_seaco = LabelSmoothingLoss( + size=self.vocab_size, + padding_idx=self.ignore_id, + smoothing=seaco_lsm_weight, + normalize_length=seaco_length_normalized_loss, + ) + self.train_decoder = kwargs.get("train_decoder", True) + self.seaco_weight = kwargs.get("seaco_weight", 0.01) + self.NO_BIAS = kwargs.get("NO_BIAS", 8377) + self.predictor_name = kwargs.get("predictor") + + # self.encoder_bmodel = EngineOV(kwargs["model_path"]+"/encoder_fp32_10b.bmodel", device_id=kwargs['dev_id']) + self.encoder_bmodel = EngineOV(kwargs["model_path"]+"/encoder_fp16_10b.bmodel", device_id=kwargs['dev_id']) + + # self.predictor_bmodel = EngineOV(kwargs["model_path"]+"/predictor_fp32_10b.bmodel", device_id=kwargs['dev_id']) + self.predictor_bmodel = EngineOV(kwargs["model_path"]+"/predictor_fp16_10b.bmodel", device_id=kwargs['dev_id']) + + self.decoder_bmodel = EngineOV(kwargs["model_path"]+"/decoder_fp32_10b.bmodel", device_id=kwargs['dev_id']) + + def forward( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + text: torch.Tensor, + text_lengths: torch.Tensor, + **kwargs, + ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]: + """Frontend + Encoder + Decoder + Calc loss + + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + text: (Batch, Length) + text_lengths: (Batch,) + """ + if len(text_lengths.size()) > 1: + text_lengths = text_lengths[:, 0] + if len(speech_lengths.size()) > 1: + speech_lengths = speech_lengths[:, 0] + # Check that batch_size is unified + assert ( + speech.shape[0] == speech_lengths.shape[0] == text.shape[0] == text_lengths.shape[0] + ), (speech.shape, speech_lengths.shape, text.shape, text_lengths.shape) + + hotword_pad = kwargs.get("hotword_pad") + hotword_lengths = kwargs.get("hotword_lengths") + seaco_label_pad = kwargs.get("seaco_label_pad") + if len(hotword_lengths.size()) > 1: + hotword_lengths = hotword_lengths[:, 0] + + batch_size = speech.shape[0] + # for data-parallel + text = text[:, : text_lengths.max()] + speech = speech[:, : speech_lengths.max()] + + # 1. Encoder + encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) + if self.predictor_bias == 1: + _, ys_pad = add_sos_eos(text, self.sos, self.eos, self.ignore_id) + ys_lengths = text_lengths + self.predictor_bias + + stats = dict() + loss_seaco = self._calc_seaco_loss( + encoder_out, + encoder_out_lens, + ys_pad, + ys_lengths, + hotword_pad, + hotword_lengths, + seaco_label_pad, + ) + if self.train_decoder: + loss_att, acc_att, _, _, _ = self._calc_att_loss( + encoder_out, encoder_out_lens, text, text_lengths + ) + loss = loss_seaco + loss_att * self.seaco_weight + stats["loss_att"] = torch.clone(loss_att.detach()) + stats["acc_att"] = acc_att + else: + loss = loss_seaco + + stats["loss_seaco"] = torch.clone(loss_seaco.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 = (text_lengths + self.predictor_bias).sum() + loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) + return loss, stats, weight + + def _merge(self, cif_attended, dec_attended): + return cif_attended + dec_attended + + def calc_predictor(self, encoder_out, encoder_out_lens): + encoder_out_mask = ( + ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :] + ).to(encoder_out.device) + predictor_outs = self.predictor( + encoder_out, None, encoder_out_mask, ignore_id=self.ignore_id + ) + return predictor_outs[:4] + + def _calc_seaco_loss( + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_lengths: torch.Tensor, + hotword_pad: torch.Tensor, + hotword_lengths: torch.Tensor, + seaco_label_pad: torch.Tensor, + ): + # predictor forward + encoder_out_mask = ( + ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :] + ).to(encoder_out.device) + pre_acoustic_embeds = self.predictor( + encoder_out, ys_pad, encoder_out_mask, ignore_id=self.ignore_id + )[0] + # decoder forward + decoder_out, _ = self.decoder( + encoder_out, encoder_out_lens, pre_acoustic_embeds, ys_lengths, return_hidden=True + ) + selected = self._hotword_representation(hotword_pad, hotword_lengths) + contextual_info = ( + selected.squeeze(0).repeat(encoder_out.shape[0], 1, 1).to(encoder_out.device) + ) + num_hot_word = contextual_info.shape[1] + _contextual_length = ( + torch.Tensor([num_hot_word]).int().repeat(encoder_out.shape[0]).to(encoder_out.device) + ) + # dha core + cif_attended, _ = self.seaco_decoder( + contextual_info, _contextual_length, pre_acoustic_embeds, ys_lengths + ) + dec_attended, _ = self.seaco_decoder( + contextual_info, _contextual_length, decoder_out, ys_lengths + ) + merged = self._merge(cif_attended, dec_attended) + dha_output = self.hotword_output_layer( + merged[:, :-1] + ) # remove the last token in loss calculation + loss_att = self.criterion_seaco(dha_output, seaco_label_pad) + return loss_att + + def _seaco_decode_with_ASF( + self, + encoder_out, + encoder_out_lens, + sematic_embeds, + ys_pad_lens, + hw_list, + nfilter=50, + seaco_weight=1.0, + ): + # decoder forward + + """ + decoder_out, decoder_hidden, _ = self.decoder( + encoder_out, + encoder_out_lens, + sematic_embeds, + ys_pad_lens, + return_hidden=True, + return_both=True, + ) + + decoder_pred = torch.log_softmax(decoder_out, dim=-1) + """ + outputs = self.decoder_bmodel([encoder_out.detach().cpu().numpy(), encoder_out_lens.detach().cpu().numpy().astype(np.int32), sematic_embeds.detach().cpu().numpy(), ys_pad_lens.detach().cpu().numpy().astype(np.int32)]) + decoder_pred, decoder_hidden = torch.from_numpy(outputs[0]), torch.from_numpy(outputs[1]) + if hw_list is not None: + hw_lengths = [len(i) for i in hw_list] + hw_list_ = [torch.Tensor(i).long() for i in hw_list] + hw_list_pad = pad_list(hw_list_, 0).to(encoder_out.device) + selected = self._hotword_representation( + hw_list_pad, torch.Tensor(hw_lengths).int().to(encoder_out.device) + ) + + contextual_info = ( + selected.squeeze(0).repeat(encoder_out.shape[0], 1, 1).to(encoder_out.device) + ) + num_hot_word = contextual_info.shape[1] + _contextual_length = ( + torch.Tensor([num_hot_word]) + .int() + .repeat(encoder_out.shape[0]) + .to(encoder_out.device) + ) + + # ASF Core + if nfilter > 0 and nfilter < num_hot_word: + hotword_scores = self.seaco_decoder.forward_asf6( + contextual_info, _contextual_length, decoder_hidden, ys_pad_lens + ) + hotword_scores = hotword_scores[0].sum(0).sum(0) + # hotword_scores /= torch.sqrt(torch.tensor(hw_lengths)[:-1].float()).to(hotword_scores.device) + dec_filter = torch.topk(hotword_scores, min(nfilter, num_hot_word - 1))[1].tolist() + add_filter = dec_filter + add_filter.append(len(hw_list_pad) - 1) + # filter hotword embedding + selected = selected[add_filter] + # again + contextual_info = ( + selected.squeeze(0).repeat(encoder_out.shape[0], 1, 1).to(encoder_out.device) + ) + num_hot_word = contextual_info.shape[1] + _contextual_length = ( + torch.Tensor([num_hot_word]) + .int() + .repeat(encoder_out.shape[0]) + .to(encoder_out.device) + ) + + # SeACo Core + cif_attended, _ = self.seaco_decoder( + contextual_info, _contextual_length, sematic_embeds, ys_pad_lens + ) + dec_attended, _ = self.seaco_decoder( + contextual_info, _contextual_length, decoder_hidden, ys_pad_lens + ) + merged = self._merge(cif_attended, dec_attended) + + dha_output = self.hotword_output_layer( + merged + ) # remove the last token in loss calculation + dha_pred = torch.log_softmax(dha_output, dim=-1) + + def _merge_res(dec_output, dha_output): + lmbd = torch.Tensor([seaco_weight] * dha_output.shape[0]) + dha_ids = dha_output.max(-1)[-1] # [0] + dha_mask = (dha_ids == self.NO_BIAS).int().unsqueeze(-1) + a = (1 - lmbd) / lmbd + b = 1 / lmbd + a, b = a.to(dec_output.device), b.to(dec_output.device) + dha_mask = (dha_mask + a.reshape(-1, 1, 1)) / b.reshape(-1, 1, 1) + # logits = dec_output * dha_mask + dha_output[:,:,:-1] * (1-dha_mask) + logits = dec_output * dha_mask + dha_output[:, :, :] * (1 - dha_mask) + return logits + + merged_pred = _merge_res(decoder_pred, dha_pred) + return merged_pred + else: + return decoder_pred + + def _hotword_representation(self, hotword_pad, hotword_lengths): + if self.bias_encoder_type != "lstm": + logging.error("Unsupported bias encoder type") + + """ + hw_embed = self.decoder.embed(hotword_pad) + hw_embed, (_, _) = self.bias_encoder(hw_embed) + if self.lstm_proj is not None: + hw_embed = self.lstm_proj(hw_embed) + _ind = np.arange(0, hw_embed.shape[0]).tolist() + selected = hw_embed[_ind, [i-1 for i in hotword_lengths.detach().cpu().tolist()]] + return selected + """ + + # hw_embed = self.sac_embedding(hotword_pad) + hw_embed = self.decoder.embed(hotword_pad) + hw_embed = torch.nn.utils.rnn.pack_padded_sequence( + hw_embed, + hotword_lengths.cpu().type(torch.int64), + batch_first=True, + enforce_sorted=False, + ) + packed_rnn_output, _ = self.bias_encoder(hw_embed) + rnn_output = torch.nn.utils.rnn.pad_packed_sequence(packed_rnn_output, batch_first=True)[0] + if self.lstm_proj is not None: + hw_hidden = self.lstm_proj(rnn_output) + else: + hw_hidden = rnn_output + _ind = np.arange(0, hw_hidden.shape[0]).tolist() + selected = hw_hidden[_ind, [i - 1 for i in hotword_lengths.detach().cpu().tolist()]] + return selected + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): + + # init beamsearch + is_use_ctc = kwargs.get("decoding_ctc_weight", 0.0) > 0.00001 and self.ctc != None + is_use_lm = ( + kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None + ) + if self.beam_search is None and (is_use_lm or is_use_ctc): + logging.info("enable beam_search") + self.init_beam_search(**kwargs) + self.nbest = kwargs.get("nbest", 1) + meta_data = {} + + # extract fbank feats + time1 = time.perf_counter() + audio_sample_list = load_audio_text_image_video( + data_in, fs=frontend.fs, audio_fs=kwargs.get("fs", 16000) + ) + time2 = time.perf_counter() + meta_data["load_data"] = f"{time2 - time1:0.3f}" + speech, speech_lengths = extract_fbank( + audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend + ) + time3 = time.perf_counter() + meta_data["extract_feat"] = f"{time3 - time2:0.3f}" + meta_data["batch_data_time"] = ( + speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 + ) + + speech = speech.to(device=kwargs["device"]) + speech_lengths = speech_lengths.to(device=kwargs["device"]) + + # hotword + """ + self.hotword_list = self.generate_hotwords_list( + kwargs.get("hotword", None), tokenizer=tokenizer, frontend=frontend + ) + """ + self.hotword_list = None + + # Encoder + #encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) + outputs = self.encoder_bmodel([speech.detach().cpu().numpy(), speech_lengths.detach().cpu().numpy()]) + encoder_out = torch.from_numpy(outputs[0]) + encoder_out_lens = speech_lengths + if isinstance(encoder_out, tuple): + encoder_out = encoder_out[0] + + """ + # predictor + predictor_outs = self.calc_predictor(encoder_out, encoder_out_lens) + pre_acoustic_embeds, pre_token_length = predictor_outs[0], predictor_outs[1] + """ + hidden, alphas, pre_token_length = torch.from_numpy(outputs[1]), torch.from_numpy(outputs[2]), torch.from_numpy(outputs[3]) + pre_acoustic_embeds, pre_peak_index = cif(hidden, alphas, 1.0) + token_num_int = torch.max(pre_token_length).type(torch.int32).item() + pre_acoustic_embeds = pre_acoustic_embeds[:, :token_num_int, :] + pre_token_length = pre_token_length.round().long() + if torch.max(pre_token_length) < 1: + return ([],) + + decoder_out = self._seaco_decode_with_ASF( + encoder_out, + encoder_out_lens, + pre_acoustic_embeds, + pre_token_length, + hw_list=self.hotword_list, + ) + + # decoder_out, _ = decoder_outs[0], decoder_outs[1] + if self.predictor_name == "CifPredictorV3": + outputs = self.predictor_bmodel([encoder_out.detach().cpu().numpy(), encoder_out_lens.detach().cpu().numpy().astype(np.int32)]) + alphas2 = torch.from_numpy(outputs[0]) + us_alphas = alphas2 * ((pre_token_length / torch.from_numpy(outputs[1]))[:, None].repeat(1, alphas2.size(1))) + us_peaks = cif_wo_hidden(us_alphas, 1.0 - 1e-4) + """ + _, _, us_alphas, us_peaks = self.calc_predictor_timestamp( + encoder_out, encoder_out_lens, pre_token_length + ) + """ + else: + us_alphas = None + + results = [] + b, n, d = decoder_out.size() + for i in range(b): + x = encoder_out[i, : encoder_out_lens[i], :] + am_scores = decoder_out[i, : pre_token_length[i], :] + if self.beam_search is not None: + nbest_hyps = self.beam_search( + x=x, + am_scores=am_scores, + maxlenratio=kwargs.get("maxlenratio", 0.0), + minlenratio=kwargs.get("minlenratio", 0.0), + ) + + nbest_hyps = nbest_hyps[: self.nbest] + else: + + yseq = am_scores.argmax(dim=-1) + score = am_scores.max(dim=-1)[0] + score = torch.sum(score, dim=-1) + # pad with mask tokens to ensure compatibility with sos/eos tokens + yseq = torch.tensor([self.sos] + yseq.tolist() + [self.eos], device=yseq.device) + nbest_hyps = [Hypothesis(yseq=yseq, score=score)] + for nbest_idx, hyp in enumerate(nbest_hyps): + ibest_writer = None + if kwargs.get("output_dir") is not None: + if not hasattr(self, "writer"): + self.writer = DatadirWriter(kwargs.get("output_dir")) + ibest_writer = self.writer[f"{nbest_idx + 1}best_recog"] + + # remove sos/eos and get results + last_pos = -1 + if isinstance(hyp.yseq, list): + token_int = hyp.yseq[1:last_pos] + else: + token_int = hyp.yseq[1:last_pos].tolist() + + # remove blank symbol id, which is assumed to be 0 + token_int = list( + filter( + lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int + ) + ) + + if tokenizer is not None: + # Change integer-ids to tokens + token = tokenizer.ids2tokens(token_int) + text = tokenizer.tokens2text(token) + if us_alphas is not None: + _, timestamp = ts_prediction_lfr6_standard( + us_alphas[i][: encoder_out_lens[i] * 3], + us_peaks[i][: encoder_out_lens[i] * 3], + copy.copy(token), + vad_offset=kwargs.get("begin_time", 0), + ) + text_postprocessed, time_stamp_postprocessed, _ = ( + postprocess_utils.sentence_postprocess(token, timestamp) + ) + result_i = { + "key": key[i], + "text": text_postprocessed, + "timestamp": time_stamp_postprocessed, + } + if ibest_writer is not None: + ibest_writer["token"][key[i]] = " ".join(token) + ibest_writer["timestamp"][key[i]] = time_stamp_postprocessed + ibest_writer["text"][key[i]] = text_postprocessed + else: + text_postprocessed, _ = postprocess_utils.sentence_postprocess(token) + result_i = {"key": key[i], "text": text_postprocessed} + if ibest_writer is not None: + ibest_writer["token"][key[i]] = " ".join(token) + ibest_writer["text"][key[i]] = text_postprocessed + else: + result_i = {"key": key[i], "token_int": token_int} + results.append(result_i) + + return results, meta_data + + def generate_hotwords_list(self, hotword_list_or_file, tokenizer=None, frontend=None): + def load_seg_dict(seg_dict_file): + seg_dict = {} + assert isinstance(seg_dict_file, str) + with open(seg_dict_file, "r", encoding="utf8") as f: + lines = f.readlines() + for line in lines: + s = line.strip().split() + key = s[0] + value = s[1:] + seg_dict[key] = " ".join(value) + return seg_dict + + def seg_tokenize(txt, seg_dict): + pattern = re.compile(r"^[\u4E00-\u9FA50-9]+$") + out_txt = "" + for word in txt: + word = word.lower() + if word in seg_dict: + out_txt += seg_dict[word] + " " + else: + if pattern.match(word): + for char in word: + if char in seg_dict: + out_txt += seg_dict[char] + " " + else: + out_txt += "" + " " + else: + out_txt += "" + " " + return out_txt.strip().split() + + seg_dict = None + if frontend.cmvn_file is not None: + model_dir = os.path.dirname(frontend.cmvn_file) + seg_dict_file = os.path.join(model_dir, "seg_dict") + if os.path.exists(seg_dict_file): + seg_dict = load_seg_dict(seg_dict_file) + else: + seg_dict = None + # for None + if hotword_list_or_file is None: + hotword_list = None + # for local txt inputs + elif os.path.exists(hotword_list_or_file) and hotword_list_or_file.endswith(".txt"): + logging.info("Attempting to parse hotwords from local txt...") + hotword_list = [] + hotword_str_list = [] + with codecs.open(hotword_list_or_file, "r") as fin: + for line in fin.readlines(): + hw = line.strip() + hw_list = hw.split() + if seg_dict is not None: + hw_list = seg_tokenize(hw_list, seg_dict) + hotword_str_list.append(hw) + hotword_list.append(tokenizer.tokens2ids(hw_list)) + hotword_list.append([self.sos]) + hotword_str_list.append("") + logging.info( + "Initialized hotword list from file: {}, hotword list: {}.".format( + hotword_list_or_file, hotword_str_list + ) + ) + # for url, download and generate txt + elif hotword_list_or_file.startswith("http"): + logging.info("Attempting to parse hotwords from url...") + work_dir = tempfile.TemporaryDirectory().name + if not os.path.exists(work_dir): + os.makedirs(work_dir) + text_file_path = os.path.join(work_dir, os.path.basename(hotword_list_or_file)) + local_file = requests.get(hotword_list_or_file) + open(text_file_path, "wb").write(local_file.content) + hotword_list_or_file = text_file_path + hotword_list = [] + hotword_str_list = [] + with codecs.open(hotword_list_or_file, "r") as fin: + for line in fin.readlines(): + hw = line.strip() + hw_list = hw.split() + if seg_dict is not None: + hw_list = seg_tokenize(hw_list, seg_dict) + hotword_str_list.append(hw) + hotword_list.append(tokenizer.tokens2ids(hw_list)) + hotword_list.append([self.sos]) + hotword_str_list.append("") + logging.info( + "Initialized hotword list from file: {}, hotword list: {}.".format( + hotword_list_or_file, hotword_str_list + ) + ) + # for text str input + elif not hotword_list_or_file.endswith(".txt"): + logging.info("Attempting to parse hotwords as str...") + hotword_list = [] + hotword_str_list = [] + for hw in hotword_list_or_file.strip().split(): + hotword_str_list.append(hw) + hw_list = hw.strip().split() + if seg_dict is not None: + hw_list = seg_tokenize(hw_list, seg_dict) + hotword_list.append(tokenizer.tokens2ids(hw_list)) + hotword_list.append([self.sos]) + hotword_str_list.append("") + logging.info("Hotword list: {}.".format(hotword_str_list)) + else: + hotword_list = None + return hotword_list + + def export( + self, + **kwargs, + ): + if "max_seq_len" not in kwargs: + kwargs["max_seq_len"] = 512 + from .export_meta import export_rebuild_model + + models = export_rebuild_model(model=self, **kwargs) + return models diff --git a/funasr/models/seaco_paraformer/template.yaml b/funasr/models/seaco_paraformer/template.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fcaf5243d449a3ced25bff7292e6e4227718df64 --- /dev/null +++ b/funasr/models/seaco_paraformer/template.yaml @@ -0,0 +1,156 @@ +# This is an example that demonstrates how to configure a model file. +# You can modify the configuration according to your own requirements. + +# to print the register_table: +# from funasr.register import tables +# tables.print() + +# network architecture +model: SeacoParaformer +model_conf: + ctc_weight: 0.0 + lsm_weight: 0.1 + length_normalized_loss: true + predictor_weight: 1.0 + predictor_bias: 1 + sampling_ratio: 0.75 + inner_dim: 512 + bias_encoder_type: lstm + bias_encoder_bid: false + seaco_lsm_weight: 0.1 + seaco_length_normal: true + train_decoder: false + NO_BIAS: 8377 + +# encoder +encoder: SANMEncoder +encoder_conf: + output_size: 512 + attention_heads: 4 + linear_units: 2048 + num_blocks: 50 + dropout_rate: 0.1 + positional_dropout_rate: 0.1 + attention_dropout_rate: 0.1 + input_layer: pe + pos_enc_class: SinusoidalPositionEncoder + normalize_before: true + kernel_size: 11 + sanm_shfit: 0 + selfattention_layer_type: sanm + +# decoder +decoder: ParaformerSANMDecoder +decoder_conf: + attention_heads: 4 + linear_units: 2048 + num_blocks: 16 + dropout_rate: 0.1 + positional_dropout_rate: 0.1 + self_attention_dropout_rate: 0.1 + src_attention_dropout_rate: 0.1 + att_layer_num: 16 + kernel_size: 11 + sanm_shfit: 0 + +# seaco decoder +seaco_decoder: ParaformerSANMDecoder +seaco_decoder_conf: + attention_heads: 4 + linear_units: 1024 + num_blocks: 4 + dropout_rate: 0.1 + positional_dropout_rate: 0.1 + self_attention_dropout_rate: 0.1 + src_attention_dropout_rate: 0.1 + kernel_size: 21 + sanm_shfit: 0 + use_output_layer: false + wo_input_layer: true + +predictor: CifPredictorV3 +predictor_conf: + idim: 512 + threshold: 1.0 + l_order: 1 + r_order: 1 + tail_threshold: 0.45 + smooth_factor2: 0.25 + noise_threshold2: 0.01 + upsample_times: 3 + use_cif1_cnn: false + upsample_type: cnn_blstm + +# frontend related +frontend: WavFrontend +frontend_conf: + fs: 16000 + window: hamming + n_mels: 80 + frame_length: 25 + frame_shift: 10 + lfr_m: 7 + lfr_n: 6 + dither: 0.0 + +specaug: SpecAugLFR +specaug_conf: + apply_time_warp: false + time_warp_window: 5 + time_warp_mode: bicubic + apply_freq_mask: true + freq_mask_width_range: + - 0 + - 30 + lfr_rate: 6 + num_freq_mask: 1 + apply_time_mask: true + time_mask_width_range: + - 0 + - 12 + num_time_mask: 1 + +train_conf: + accum_grad: 1 + grad_clip: 5 + max_epoch: 150 + val_scheduler_criterion: + - valid + - acc + best_model_criterion: + - - valid + - acc + - max + keep_nbest_models: 10 + log_interval: 50 + +optim: adam +optim_conf: + lr: 0.0005 +scheduler: warmuplr +scheduler_conf: + warmup_steps: 30000 + +dataset: AudioDataset +dataset_conf: + index_ds: IndexDSJsonl + batch_sampler: BatchSampler + batch_type: example # example or length + batch_size: 1 # if batch_type is example, batch_size is the numbers of samples; if length, batch_size is source_token_len+target_token_len; + max_token_length: 2048 # filter samples if source_token_len+target_token_len > max_token_length, + buffer_size: 500 + shuffle: True + num_workers: 0 + +tokenizer: CharTokenizer +tokenizer_conf: + unk_symbol: + split_with_space: true + + +ctc_conf: + dropout_rate: 0.0 + ctc_type: builtin + reduce: true + ignore_nan_grad: true +normalize: null diff --git a/funasr/models/sense_voice/__init__.py b/funasr/models/sense_voice/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/models/sense_voice/export_meta.py b/funasr/models/sense_voice/export_meta.py new file mode 100644 index 0000000000000000000000000000000000000000..449388ef79e8fec4f0a33287b68db9309c2c0037 --- /dev/null +++ b/funasr/models/sense_voice/export_meta.py @@ -0,0 +1,77 @@ +#!/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 types +import torch +from funasr.utils.torch_function import sequence_mask + + +def export_rebuild_model(model, **kwargs): + model.device = kwargs.get("device") + model.make_pad_mask = sequence_mask(kwargs["max_seq_len"], flip=False) + model.forward = types.MethodType(export_forward, model) + model.export_dummy_inputs = types.MethodType(export_dummy_inputs, model) + model.export_input_names = types.MethodType(export_input_names, model) + model.export_output_names = types.MethodType(export_output_names, model) + model.export_dynamic_axes = types.MethodType(export_dynamic_axes, model) + model.export_name = types.MethodType(export_name, model) + return model + +def export_forward( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + language: torch.Tensor, + textnorm: torch.Tensor, + **kwargs, +): + # speech = speech.to(device="cuda") + # speech_lengths = speech_lengths.to(device="cuda") + language_query = self.embed(language.to(speech.device)).unsqueeze(1) + textnorm_query = self.embed(textnorm.to(speech.device)).unsqueeze(1) + + speech = torch.cat((textnorm_query, speech), dim=1) + + event_emo_query = self.embed(torch.LongTensor([[1, 2]]).to(speech.device)).repeat( + speech.size(0), 1, 1 + ) + input_query = torch.cat((language_query, event_emo_query), dim=1) + speech = torch.cat((input_query, speech), dim=1) + + speech_lengths_new = speech_lengths + 4 + encoder_out, encoder_out_lens = self.encoder(speech, speech_lengths_new) + + if isinstance(encoder_out, tuple): + encoder_out = encoder_out[0] + + ctc_logits = self.ctc.ctc_lo(encoder_out) + + return ctc_logits, encoder_out_lens + +def export_dummy_inputs(self): + speech = torch.randn(2, 30, 560) + speech_lengths = torch.tensor([6, 30], dtype=torch.int32) + language = torch.tensor([0, 0], dtype=torch.int32) + textnorm = torch.tensor([15, 15], dtype=torch.int32) + return (speech, speech_lengths, language, textnorm) + +def export_input_names(self): + return ["speech", "speech_lengths", "language", "textnorm"] + +def export_output_names(self): + return ["ctc_logits", "encoder_out_lens"] + +def export_dynamic_axes(self): + return { + "speech": {0: "batch_size", 1: "feats_length"}, + "speech_lengths": {0: "batch_size"}, + "language": {0: "batch_size"}, + "textnorm": {0: "batch_size"}, + "ctc_logits": {0: "batch_size", 1: "logits_length"}, + "encoder_out_lens": {0: "batch_size"}, + } + +def export_name(self): + return "model.onnx" diff --git a/funasr/models/sense_voice/model.py b/funasr/models/sense_voice/model.py new file mode 100644 index 0000000000000000000000000000000000000000..bd1c76dbc9f3bd78174d403a32a393c8e238cb7d --- /dev/null +++ b/funasr/models/sense_voice/model.py @@ -0,0 +1,990 @@ +from typing import Iterable, Optional +import types +import time +import numpy as np +import torch +import torch.nn.functional as F +from torch import Tensor +from torch import nn +from torch.cuda.amp import autocast +from funasr.metrics.compute_acc import compute_accuracy, th_accuracy +from funasr.losses.label_smoothing_loss import LabelSmoothingLoss +from funasr.train_utils.device_funcs import force_gatherable + +from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank +from funasr.utils.datadir_writer import DatadirWriter +from funasr.models.ctc.ctc import CTC + +from funasr.register import tables + + +from funasr.models.paraformer.search import Hypothesis +from .utils.ctc_alignment import ctc_forced_align + + +class SinusoidalPositionEncoder(torch.nn.Module): + """ """ + + def __int__(self, d_model=80, dropout_rate=0.1): + pass + + def encode( + self, positions: torch.Tensor = None, depth: int = None, dtype: torch.dtype = torch.float32 + ): + batch_size = positions.size(0) + positions = positions.type(dtype) + device = positions.device + log_timescale_increment = torch.log(torch.tensor([10000], dtype=dtype, device=device)) / ( + depth / 2 - 1 + ) + inv_timescales = torch.exp( + torch.arange(depth / 2, device=device).type(dtype) * (-log_timescale_increment) + ) + inv_timescales = torch.reshape(inv_timescales, [batch_size, -1]) + scaled_time = torch.reshape(positions, [1, -1, 1]) * torch.reshape( + inv_timescales, [1, 1, -1] + ) + encoding = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=2) + return encoding.type(dtype) + + def forward(self, x): + batch_size, timesteps, input_dim = x.size() + positions = torch.arange(1, timesteps + 1, device=x.device)[None, :] + position_encoding = self.encode(positions, input_dim, x.dtype).to(x.device) + + return x + position_encoding + + +class PositionwiseFeedForward(torch.nn.Module): + """Positionwise feed forward layer. + + Args: + idim (int): Input dimenstion. + hidden_units (int): The number of hidden units. + dropout_rate (float): Dropout rate. + + """ + + def __init__(self, idim, hidden_units, dropout_rate, activation=torch.nn.ReLU()): + """Construct an PositionwiseFeedForward object.""" + super(PositionwiseFeedForward, self).__init__() + self.w_1 = torch.nn.Linear(idim, hidden_units) + self.w_2 = torch.nn.Linear(hidden_units, idim) + self.dropout = torch.nn.Dropout(dropout_rate) + self.activation = activation + + def forward(self, x): + """Forward function.""" + return self.w_2(self.dropout(self.activation(self.w_1(x)))) + + +class MultiHeadedAttentionSANM(nn.Module): + """Multi-Head Attention layer. + + Args: + n_head (int): The number of heads. + n_feat (int): The number of features. + dropout_rate (float): Dropout rate. + + """ + + def __init__( + self, + n_head, + in_feat, + n_feat, + dropout_rate, + kernel_size, + sanm_shfit=0, + lora_list=None, + lora_rank=8, + lora_alpha=16, + lora_dropout=0.1, + ): + """Construct an MultiHeadedAttention object.""" + super().__init__() + assert n_feat % n_head == 0 + # We assume d_v always equals d_k + self.d_k = n_feat // n_head + self.h = n_head + # self.linear_q = nn.Linear(n_feat, n_feat) + # self.linear_k = nn.Linear(n_feat, n_feat) + # self.linear_v = nn.Linear(n_feat, n_feat) + + self.linear_out = nn.Linear(n_feat, n_feat) + self.linear_q_k_v = nn.Linear(in_feat, n_feat * 3) + self.attn = None + self.dropout = nn.Dropout(p=dropout_rate) + + self.fsmn_block = nn.Conv1d( + n_feat, n_feat, kernel_size, stride=1, padding=0, groups=n_feat, bias=False + ) + # padding + left_padding = (kernel_size - 1) // 2 + if sanm_shfit > 0: + left_padding = left_padding + sanm_shfit + right_padding = kernel_size - 1 - left_padding + self.pad_fn = nn.ConstantPad1d((left_padding, right_padding), 0.0) + + def forward_fsmn(self, inputs, mask, mask_shfit_chunk=None): + b, t, d = inputs.size() + if mask is not None: + mask = torch.reshape(mask, (b, -1, 1)) + if mask_shfit_chunk is not None: + mask = mask * mask_shfit_chunk + inputs = inputs * mask + + x = inputs.transpose(1, 2) + x = self.pad_fn(x) + x = self.fsmn_block(x) + x = x.transpose(1, 2) + x += inputs + x = self.dropout(x) + if mask is not None: + x = x * mask + return x + + def forward_qkv(self, x): + """Transform query, key and value. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + + Returns: + torch.Tensor: Transformed query tensor (#batch, n_head, time1, d_k). + torch.Tensor: Transformed key tensor (#batch, n_head, time2, d_k). + torch.Tensor: Transformed value tensor (#batch, n_head, time2, d_k). + + """ + b, t, d = x.size() + q_k_v = self.linear_q_k_v(x) + q, k, v = torch.split(q_k_v, int(self.h * self.d_k), dim=-1) + q_h = torch.reshape(q, (b, t, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time1, d_k) + k_h = torch.reshape(k, (b, t, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time2, d_k) + v_h = torch.reshape(v, (b, t, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time2, d_k) + + return q_h, k_h, v_h, v + + def forward_attention(self, value, scores, mask, mask_att_chunk_encoder=None): + """Compute attention context vector. + + Args: + value (torch.Tensor): Transformed value (#batch, n_head, time2, d_k). + scores (torch.Tensor): Attention score (#batch, n_head, time1, time2). + mask (torch.Tensor): Mask (#batch, 1, time2) or (#batch, time1, time2). + + Returns: + torch.Tensor: Transformed value (#batch, time1, d_model) + weighted by the attention score (#batch, time1, time2). + + """ + n_batch = value.size(0) + if mask is not None: + if mask_att_chunk_encoder is not None: + mask = mask * mask_att_chunk_encoder + + mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) + + min_value = -float( + "inf" + ) # float(numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min) + scores = scores.masked_fill(mask, min_value) + attn = torch.softmax(scores, dim=-1).masked_fill( + mask, 0.0 + ) # (batch, head, time1, time2) + else: + attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2) + + p_attn = self.dropout(attn) + x = torch.matmul(p_attn, value) # (batch, head, time1, d_k) + x = ( + x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) + ) # (batch, time1, d_model) + + return self.linear_out(x) # (batch, time1, d_model) + + def forward(self, x, mask, mask_shfit_chunk=None, mask_att_chunk_encoder=None): + """Compute scaled dot product attention. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + mask (torch.Tensor): Mask tensor (#batch, 1, time2) or + (#batch, time1, time2). + + Returns: + torch.Tensor: Output tensor (#batch, time1, d_model). + + """ + q_h, k_h, v_h, v = self.forward_qkv(x) + fsmn_memory = self.forward_fsmn(v, mask, mask_shfit_chunk) + q_h = q_h * self.d_k ** (-0.5) + scores = torch.matmul(q_h, k_h.transpose(-2, -1)) + att_outs = self.forward_attention(v_h, scores, mask, mask_att_chunk_encoder) + return att_outs + fsmn_memory + + def forward_chunk(self, x, cache=None, chunk_size=None, look_back=0): + """Compute scaled dot product attention. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + mask (torch.Tensor): Mask tensor (#batch, 1, time2) or + (#batch, time1, time2). + + Returns: + torch.Tensor: Output tensor (#batch, time1, d_model). + + """ + q_h, k_h, v_h, v = self.forward_qkv(x) + if chunk_size is not None and look_back > 0 or look_back == -1: + if cache is not None: + k_h_stride = k_h[:, :, : -(chunk_size[2]), :] + v_h_stride = v_h[:, :, : -(chunk_size[2]), :] + k_h = torch.cat((cache["k"], k_h), dim=2) + v_h = torch.cat((cache["v"], v_h), dim=2) + + cache["k"] = torch.cat((cache["k"], k_h_stride), dim=2) + cache["v"] = torch.cat((cache["v"], v_h_stride), dim=2) + if look_back != -1: + cache["k"] = cache["k"][:, :, -(look_back * chunk_size[1]) :, :] + cache["v"] = cache["v"][:, :, -(look_back * chunk_size[1]) :, :] + else: + cache_tmp = { + "k": k_h[:, :, : -(chunk_size[2]), :], + "v": v_h[:, :, : -(chunk_size[2]), :], + } + cache = cache_tmp + fsmn_memory = self.forward_fsmn(v, None) + q_h = q_h * self.d_k ** (-0.5) + scores = torch.matmul(q_h, k_h.transpose(-2, -1)) + att_outs = self.forward_attention(v_h, scores, None) + return att_outs + fsmn_memory, cache + + +class LayerNorm(nn.LayerNorm): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward(self, input): + output = F.layer_norm( + input.float(), + self.normalized_shape, + self.weight.float() if self.weight is not None else None, + self.bias.float() if self.bias is not None else None, + self.eps, + ) + return output.type_as(input) + + +def sequence_mask(lengths, maxlen=None, dtype=torch.float32, device=None): + if maxlen is None: + maxlen = lengths.max() + row_vector = torch.arange(0, maxlen, 1).to(lengths.device) + matrix = torch.unsqueeze(lengths, dim=-1) + mask = row_vector < matrix + mask = mask.detach() + + return mask.type(dtype).to(device) if device is not None else mask.type(dtype) + + +class EncoderLayerSANM(nn.Module): + def __init__( + self, + in_size, + size, + self_attn, + feed_forward, + dropout_rate, + normalize_before=True, + concat_after=False, + stochastic_depth_rate=0.0, + ): + """Construct an EncoderLayer object.""" + super(EncoderLayerSANM, self).__init__() + self.self_attn = self_attn + self.feed_forward = feed_forward + self.norm1 = LayerNorm(in_size) + self.norm2 = LayerNorm(size) + self.dropout = nn.Dropout(dropout_rate) + self.in_size = in_size + self.size = size + self.normalize_before = normalize_before + self.concat_after = concat_after + if self.concat_after: + self.concat_linear = nn.Linear(size + size, size) + self.stochastic_depth_rate = stochastic_depth_rate + self.dropout_rate = dropout_rate + + def forward(self, x, mask, cache=None, mask_shfit_chunk=None, mask_att_chunk_encoder=None): + """Compute encoded features. + + Args: + x_input (torch.Tensor): Input tensor (#batch, time, size). + mask (torch.Tensor): Mask tensor for the input (#batch, time). + cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size). + + Returns: + torch.Tensor: Output tensor (#batch, time, size). + torch.Tensor: Mask tensor (#batch, time). + + """ + skip_layer = False + # with stochastic depth, residual connection `x + f(x)` becomes + # `x <- x + 1 / (1 - p) * f(x)` at training time. + stoch_layer_coeff = 1.0 + if self.training and self.stochastic_depth_rate > 0: + skip_layer = torch.rand(1).item() < self.stochastic_depth_rate + stoch_layer_coeff = 1.0 / (1 - self.stochastic_depth_rate) + + if skip_layer: + if cache is not None: + x = torch.cat([cache, x], dim=1) + return x, mask + + residual = x + if self.normalize_before: + x = self.norm1(x) + + if self.concat_after: + x_concat = torch.cat( + ( + x, + self.self_attn( + x, + mask, + mask_shfit_chunk=mask_shfit_chunk, + mask_att_chunk_encoder=mask_att_chunk_encoder, + ), + ), + dim=-1, + ) + if self.in_size == self.size: + x = residual + stoch_layer_coeff * self.concat_linear(x_concat) + else: + x = stoch_layer_coeff * self.concat_linear(x_concat) + else: + if self.in_size == self.size: + x = residual + stoch_layer_coeff * self.dropout( + self.self_attn( + x, + mask, + mask_shfit_chunk=mask_shfit_chunk, + mask_att_chunk_encoder=mask_att_chunk_encoder, + ) + ) + else: + x = stoch_layer_coeff * self.dropout( + self.self_attn( + x, + mask, + mask_shfit_chunk=mask_shfit_chunk, + mask_att_chunk_encoder=mask_att_chunk_encoder, + ) + ) + if not self.normalize_before: + x = self.norm1(x) + + residual = x + if self.normalize_before: + x = self.norm2(x) + x = residual + stoch_layer_coeff * self.dropout(self.feed_forward(x)) + if not self.normalize_before: + x = self.norm2(x) + + return x, mask, cache, mask_shfit_chunk, mask_att_chunk_encoder + + def forward_chunk(self, x, cache=None, chunk_size=None, look_back=0): + """Compute encoded features. + + Args: + x_input (torch.Tensor): Input tensor (#batch, time, size). + mask (torch.Tensor): Mask tensor for the input (#batch, time). + cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size). + + Returns: + torch.Tensor: Output tensor (#batch, time, size). + torch.Tensor: Mask tensor (#batch, time). + + """ + + residual = x + if self.normalize_before: + x = self.norm1(x) + + if self.in_size == self.size: + attn, cache = self.self_attn.forward_chunk(x, cache, chunk_size, look_back) + x = residual + attn + else: + x, cache = self.self_attn.forward_chunk(x, cache, chunk_size, look_back) + + if not self.normalize_before: + x = self.norm1(x) + + residual = x + if self.normalize_before: + x = self.norm2(x) + x = residual + self.feed_forward(x) + if not self.normalize_before: + x = self.norm2(x) + + return x, cache + + +@tables.register("encoder_classes", "SenseVoiceEncoderSmall") +class SenseVoiceEncoderSmall(nn.Module): + """ + Author: Speech Lab of DAMO Academy, Alibaba Group + SCAMA: Streaming chunk-aware multihead attention for online end-to-end speech recognition + https://arxiv.org/abs/2006.01713 + """ + + def __init__( + self, + input_size: int, + output_size: int = 256, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + tp_blocks: int = 0, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + attention_dropout_rate: float = 0.0, + stochastic_depth_rate: float = 0.0, + input_layer: Optional[str] = "conv2d", + pos_enc_class=SinusoidalPositionEncoder, + normalize_before: bool = True, + concat_after: bool = False, + positionwise_layer_type: str = "linear", + positionwise_conv_kernel_size: int = 1, + padding_idx: int = -1, + kernel_size: int = 11, + sanm_shfit: int = 0, + selfattention_layer_type: str = "sanm", + **kwargs, + ): + super().__init__() + self._output_size = output_size + + self.embed = SinusoidalPositionEncoder() + + self.normalize_before = normalize_before + + positionwise_layer = PositionwiseFeedForward + positionwise_layer_args = ( + output_size, + linear_units, + dropout_rate, + ) + + encoder_selfattn_layer = MultiHeadedAttentionSANM + encoder_selfattn_layer_args0 = ( + attention_heads, + input_size, + output_size, + attention_dropout_rate, + kernel_size, + sanm_shfit, + ) + encoder_selfattn_layer_args = ( + attention_heads, + output_size, + output_size, + attention_dropout_rate, + kernel_size, + sanm_shfit, + ) + + self.encoders0 = nn.ModuleList( + [ + EncoderLayerSANM( + input_size, + output_size, + encoder_selfattn_layer(*encoder_selfattn_layer_args0), + positionwise_layer(*positionwise_layer_args), + dropout_rate, + ) + for i in range(1) + ] + ) + self.encoders = nn.ModuleList( + [ + EncoderLayerSANM( + output_size, + output_size, + encoder_selfattn_layer(*encoder_selfattn_layer_args), + positionwise_layer(*positionwise_layer_args), + dropout_rate, + ) + for i in range(num_blocks - 1) + ] + ) + + self.tp_encoders = nn.ModuleList( + [ + EncoderLayerSANM( + output_size, + output_size, + encoder_selfattn_layer(*encoder_selfattn_layer_args), + positionwise_layer(*positionwise_layer_args), + dropout_rate, + ) + for i in range(tp_blocks) + ] + ) + + self.after_norm = LayerNorm(output_size) + + self.tp_norm = LayerNorm(output_size) + + def output_size(self) -> int: + return self._output_size + + def forward( + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + ): + """Embed positions in tensor.""" + maxlen = xs_pad.shape[1] + masks = sequence_mask(ilens, maxlen=maxlen, device=ilens.device)[:, None, :] + + xs_pad *= self.output_size() ** 0.5 + + xs_pad = self.embed(xs_pad) + + # forward encoder1 + for layer_idx, encoder_layer in enumerate(self.encoders0): + encoder_outs = encoder_layer(xs_pad, masks) + xs_pad, masks = encoder_outs[0], encoder_outs[1] + + for layer_idx, encoder_layer in enumerate(self.encoders): + encoder_outs = encoder_layer(xs_pad, masks) + xs_pad, masks = encoder_outs[0], encoder_outs[1] + + xs_pad = self.after_norm(xs_pad) + + # forward encoder2 + olens = masks.squeeze(1).sum(1).int() + + for layer_idx, encoder_layer in enumerate(self.tp_encoders): + encoder_outs = encoder_layer(xs_pad, masks) + xs_pad, masks = encoder_outs[0], encoder_outs[1] + + xs_pad = self.tp_norm(xs_pad) + return xs_pad, olens + + +@tables.register("model_classes", "SenseVoiceSmall") +class SenseVoiceSmall(nn.Module): + """CTC-attention hybrid Encoder-Decoder model""" + + def __init__( + self, + specaug: str = None, + specaug_conf: dict = None, + normalize: str = None, + normalize_conf: dict = None, + encoder: str = None, + encoder_conf: dict = None, + ctc_conf: dict = None, + input_size: int = 80, + vocab_size: int = -1, + ignore_id: int = -1, + blank_id: int = 0, + sos: int = 1, + eos: int = 2, + length_normalized_loss: 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() + + if ctc_conf is None: + ctc_conf = {} + ctc = CTC(odim=vocab_size, encoder_output_size=encoder_output_size, **ctc_conf) + + self.blank_id = blank_id + self.sos = sos if sos is not None else vocab_size - 1 + self.eos = eos if eos is not None else vocab_size - 1 + self.vocab_size = vocab_size + self.ignore_id = ignore_id + self.specaug = specaug + self.normalize = normalize + self.encoder = encoder + self.error_calculator = None + + self.ctc = ctc + + self.length_normalized_loss = length_normalized_loss + self.encoder_output_size = encoder_output_size + + self.lid_dict = {"auto": 0, "zh": 3, "en": 4, "yue": 7, "ja": 11, "ko": 12, "nospeech": 13} + self.lid_int_dict = {24884: 3, 24885: 4, 24888: 7, 24892: 11, 24896: 12, 24992: 13} + self.textnorm_dict = {"withitn": 14, "woitn": 15} + self.textnorm_int_dict = {25016: 14, 25017: 15} + self.embed = torch.nn.Embedding( + 7 + len(self.lid_dict) + len(self.textnorm_dict), input_size + ) + self.emo_dict = { + "unk": 25009, + "happy": 25001, + "sad": 25002, + "angry": 25003, + "neutral": 25004, + } + + self.criterion_att = LabelSmoothingLoss( + size=self.vocab_size, + padding_idx=self.ignore_id, + smoothing=kwargs.get("lsm_weight", 0.0), + normalize_length=self.length_normalized_loss, + ) + + @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 + + def forward( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + text: torch.Tensor, + text_lengths: torch.Tensor, + **kwargs, + ): + """Encoder + Decoder + Calc loss + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + text: (Batch, Length) + text_lengths: (Batch,) + """ + # import pdb; + # pdb.set_trace() + 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] + + # 1. Encoder + encoder_out, encoder_out_lens = self.encode(speech, speech_lengths, text) + + loss_ctc, cer_ctc = None, None + loss_rich, acc_rich = None, None + stats = dict() + + loss_ctc, cer_ctc = self._calc_ctc_loss( + encoder_out[:, 4:, :], encoder_out_lens - 4, text[:, 4:], text_lengths - 4 + ) + + loss_rich, acc_rich = self._calc_rich_ce_loss(encoder_out[:, :4, :], text[:, :4]) + + loss = loss_ctc + loss_rich + # Collect total loss stats + stats["loss_ctc"] = torch.clone(loss_ctc.detach()) if loss_ctc is not None else None + stats["loss_rich"] = torch.clone(loss_rich.detach()) if loss_rich is not None else None + stats["loss"] = torch.clone(loss.detach()) if loss is not None else None + stats["acc_rich"] = acc_rich + + # 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 encode( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + text: torch.Tensor, + **kwargs, + ): + """Frontend + Encoder. Note that this method is used by asr_inference.py + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + ind: int + """ + + # 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) + + lids = torch.LongTensor( + [ + [ + ( + self.lid_int_dict[int(lid)] + if torch.rand(1) > 0.2 and int(lid) in self.lid_int_dict + else 0 + ) + ] + for lid in text[:, 0] + ] + ).to(speech.device) + language_query = self.embed(lids) + + styles = torch.LongTensor( + [[self.textnorm_int_dict[int(style)]] for style in text[:, 3]] + ).to(speech.device) + style_query = self.embed(styles) + speech = torch.cat((style_query, speech), dim=1) + speech_lengths += 1 + + event_emo_query = self.embed(torch.LongTensor([[1, 2]]).to(speech.device)).repeat( + speech.size(0), 1, 1 + ) + input_query = torch.cat((language_query, event_emo_query), dim=1) + speech = torch.cat((input_query, speech), dim=1) + speech_lengths += 3 + + encoder_out, encoder_out_lens = self.encoder(speech, speech_lengths) + + return encoder_out, encoder_out_lens + + def _calc_ctc_loss( + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, + ): + # Calc CTC loss + loss_ctc = self.ctc(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens) + + # Calc CER using CTC + cer_ctc = None + if not self.training and self.error_calculator is not None: + ys_hat = self.ctc.argmax(encoder_out).data + cer_ctc = self.error_calculator(ys_hat.cpu(), ys_pad.cpu(), is_ctc=True) + return loss_ctc, cer_ctc + + def _calc_rich_ce_loss( + self, + encoder_out: torch.Tensor, + ys_pad: torch.Tensor, + ): + decoder_out = self.ctc.ctc_lo(encoder_out) + # 2. Compute attention loss + loss_rich = self.criterion_att(decoder_out, ys_pad.contiguous()) + acc_rich = th_accuracy( + decoder_out.view(-1, self.vocab_size), + ys_pad.contiguous(), + ignore_label=self.ignore_id, + ) + + return loss_rich, acc_rich + + def inference( + self, + data_in, + data_lengths=None, + key: list = ["wav_file_tmp_name"], + tokenizer=None, + frontend=None, + **kwargs, + ): + + 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"]) + + language = kwargs.get("language", "auto") + language_query = self.embed( + torch.LongTensor([[self.lid_dict[language] if language in self.lid_dict else 0]]).to( + speech.device + ) + ).repeat(speech.size(0), 1, 1) + + use_itn = kwargs.get("use_itn", False) + textnorm = kwargs.get("text_norm", None) + output_timestamp = kwargs.get("output_timestamp", False) + + if textnorm is None: + textnorm = "withitn" if use_itn else "woitn" + textnorm_query = self.embed( + torch.LongTensor([[self.textnorm_dict[textnorm]]]).to(speech.device) + ).repeat(speech.size(0), 1, 1) + speech = torch.cat((textnorm_query, speech), dim=1) + speech_lengths += 1 + + event_emo_query = self.embed(torch.LongTensor([[1, 2]]).to(speech.device)).repeat( + speech.size(0), 1, 1 + ) + input_query = torch.cat((language_query, event_emo_query), dim=1) + speech = torch.cat((input_query, speech), dim=1) + speech_lengths += 3 + + # Encoder + encoder_out, encoder_out_lens = self.encoder(speech, speech_lengths) + if isinstance(encoder_out, tuple): + encoder_out = encoder_out[0] + + # c. Passed the encoder result and the beam search + ctc_logits = self.ctc.log_softmax(encoder_out) + if kwargs.get("ban_emo_unk", False): + ctc_logits[:, :, self.emo_dict["unk"]] = -float("inf") + + results = [] + 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) + + 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"1best_recog"] + + mask = yseq != self.blank_id + token_int = yseq[mask].tolist() + + # Change integer-ids to tokens + text = tokenizer.decode(token_int) + + # result_i = {"key": key[i], "text": text} + # results.append(result_i) + + if ibest_writer is not None: + ibest_writer["text"][key[i]] = text + + if output_timestamp: + from itertools import groupby + + timestamp = [] + tokens = tokenizer.text2tokens(text)[4:] + token_back_to_id = tokenizer.tokens2ids(tokens) + token_ids = [] + for tok_ls in token_back_to_id: + if tok_ls: token_ids.extend(tok_ls) + else: token_ids.append(124) + + logits_speech = self.ctc.softmax(encoder_out)[i, 4 : encoder_out_lens[i].item(), :] + pred = logits_speech.argmax(-1).cpu() + logits_speech[pred == self.blank_id, self.blank_id] = 0 + align = ctc_forced_align( + logits_speech.unsqueeze(0).float(), + torch.Tensor(token_ids).unsqueeze(0).long().to(logits_speech.device), + (encoder_out_lens[i] - 4).long(), + torch.tensor(len(token_ids)).unsqueeze(0).long().to(logits_speech.device), + ignore_id=self.ignore_id, + ) + pred = groupby(align[0, : encoder_out_lens[i]]) + _start = 0 + token_id = 0 + ts_max = encoder_out_lens[i] - 4 + for pred_token, pred_frame in pred: + _end = _start + len(list(pred_frame)) + if pred_token != 0: + ts_left = max((_start * 60 - 30) / 1000, 0) + ts_right = min((_end * 60 - 30) / 1000, (ts_max * 60 - 30) / 1000) + timestamp.append([tokens[token_id], ts_left, ts_right]) + token_id += 1 + _start = _end + timestamp = self.post(timestamp) + result_i = {"key": key[i], "text": text, "timestamp": timestamp} + results.append(result_i) + else: + result_i = {"key": key[i], "text": text} + results.append(result_i) + return results, meta_data + + def post(self, timestamp): + timestamp_new = [] + prev_word = None + for i, t in enumerate(timestamp): + word, start, end = t + start = int(start * 1000) + end = int(end * 1000) + if word == "▁": + continue + if i == 0: + # timestamp_new.append([word, start, end]) + timestamp_new.append([start, end]) + elif word.startswith("▁"): + word = word[1:] + timestamp_new.append([start, end]) + elif prev_word is not None and prev_word.isalpha() and prev_word.isascii() and word.isalpha() and word.isascii(): + prev_word += word + timestamp_new[-1][1] = end + else: + # timestamp_new[-1][0] += word + timestamp_new.append([start, end]) + prev_word = word + return timestamp_new + + def export(self, **kwargs): + from .export_meta import export_rebuild_model + + if "max_seq_len" not in kwargs: + kwargs["max_seq_len"] = 512 + models = export_rebuild_model(model=self, **kwargs) + return models + + return results, meta_data diff --git a/funasr/models/sense_voice/utils/__init__.py b/funasr/models/sense_voice/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/models/sense_voice/utils/ctc_alignment.py b/funasr/models/sense_voice/utils/ctc_alignment.py new file mode 100644 index 0000000000000000000000000000000000000000..e694e2094e8db1a7db8e0fd7dd6a1436bff4a214 --- /dev/null +++ b/funasr/models/sense_voice/utils/ctc_alignment.py @@ -0,0 +1,65 @@ +import torch +def ctc_forced_align( + log_probs: torch.Tensor, + targets: torch.Tensor, + input_lengths: torch.Tensor, + target_lengths: torch.Tensor, + blank: int = 0, + ignore_id: int = -1, +) -> torch.Tensor: + """Align a CTC label sequence to an emission. + Args: + log_probs (Tensor): log probability of CTC emission output. + Tensor of shape `(B, T, C)`. where `B` is the batch size, `T` is the input length, + `C` is the number of characters in alphabet including blank. + targets (Tensor): Target sequence. Tensor of shape `(B, L)`, + where `L` is the target length. + input_lengths (Tensor): + Lengths of the inputs (max value must each be <= `T`). 1-D Tensor of shape `(B,)`. + target_lengths (Tensor): + Lengths of the targets. 1-D Tensor of shape `(B,)`. + blank_id (int, optional): The index of blank symbol in CTC emission. (Default: 0) + ignore_id (int, optional): The index of ignore symbol in CTC emission. (Default: -1) + """ + targets[targets == ignore_id] = blank + batch_size, input_time_size, _ = log_probs.size() + bsz_indices = torch.arange(batch_size, device=input_lengths.device) + _t_a_r_g_e_t_s_ = torch.cat( + ( + torch.stack((torch.full_like(targets, blank), targets), dim=-1).flatten(start_dim=1), + torch.full_like(targets[:, :1], blank), + ), + dim=-1, + ) + diff_labels = torch.cat( + ( + torch.as_tensor([[False, False]], device=targets.device).expand(batch_size, -1), + _t_a_r_g_e_t_s_[:, 2:] != _t_a_r_g_e_t_s_[:, :-2], + ), + dim=1, + ) + neg_inf = torch.tensor(float("-inf"), device=log_probs.device, dtype=log_probs.dtype) + padding_num = 2 + padded_t = padding_num + _t_a_r_g_e_t_s_.size(-1) + best_score = torch.full((batch_size, padded_t), neg_inf, device=log_probs.device, dtype=log_probs.dtype) + best_score[:, padding_num + 0] = log_probs[:, 0, blank] + best_score[:, padding_num + 1] = log_probs[bsz_indices, 0, _t_a_r_g_e_t_s_[:, 1]] + backpointers = torch.zeros((batch_size, input_time_size, padded_t), device=log_probs.device, dtype=targets.dtype) + for t in range(1, input_time_size): + prev = torch.stack( + (best_score[:, 2:], best_score[:, 1:-1], torch.where(diff_labels, best_score[:, :-2], neg_inf)) + ) + prev_max_value, prev_max_idx = prev.max(dim=0) + best_score[:, padding_num:] = log_probs[:, t].gather(-1, _t_a_r_g_e_t_s_) + prev_max_value + backpointers[:, t, padding_num:] = prev_max_idx + l1l2 = best_score.gather( + -1, torch.stack((padding_num + target_lengths * 2 - 1, padding_num + target_lengths * 2), dim=-1) + ) + path = torch.zeros((batch_size, input_time_size), device=best_score.device, dtype=torch.long) + path[bsz_indices, input_lengths - 1] = padding_num + target_lengths * 2 - 1 + l1l2.argmax(dim=-1) + for t in range(input_time_size - 1, 0, -1): + target_indices = path[:, t] + prev_max_idx = backpointers[bsz_indices, t, target_indices] + path[:, t - 1] += target_indices - prev_max_idx + alignments = _t_a_r_g_e_t_s_.gather(dim=-1, index=(path - padding_num).clamp(min=0)) + return alignments \ No newline at end of file diff --git a/funasr/models/sond/__init__.py b/funasr/models/sond/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/models/sond/attention.py b/funasr/models/sond/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..18580b7fe4084fb8f6437d523660f9bc87879da9 --- /dev/null +++ b/funasr/models/sond/attention.py @@ -0,0 +1,322 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +# Copyright 2019 Shigeki Karita +# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +"""Multi-Head Attention layer definition.""" + +import math + +import numpy +import torch +from torch import nn +from typing import Optional, Tuple + +import torch.nn.functional as F +from funasr.models.transformer.utils.nets_utils import make_pad_mask +import funasr.models.lora.layers as lora + + +class MultiHeadedAttention(nn.Module): + """Multi-Head Attention layer. + + Args: + n_head (int): The number of heads. + n_feat (int): The number of features. + dropout_rate (float): Dropout rate. + + """ + + def __init__(self, n_head, n_feat, dropout_rate): + """Construct an MultiHeadedAttention object.""" + super(MultiHeadedAttention, self).__init__() + assert n_feat % n_head == 0 + # We assume d_v always equals d_k + self.d_k = n_feat // n_head + self.h = n_head + self.linear_q = nn.Linear(n_feat, n_feat) + self.linear_k = nn.Linear(n_feat, n_feat) + self.linear_v = nn.Linear(n_feat, n_feat) + self.linear_out = nn.Linear(n_feat, n_feat) + self.attn = None + self.dropout = nn.Dropout(p=dropout_rate) + + def forward_qkv(self, query, key, value): + """Transform query, key and value. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + + Returns: + torch.Tensor: Transformed query tensor (#batch, n_head, time1, d_k). + torch.Tensor: Transformed key tensor (#batch, n_head, time2, d_k). + torch.Tensor: Transformed value tensor (#batch, n_head, time2, d_k). + + """ + n_batch = query.size(0) + q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k) + k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k) + v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k) + q = q.transpose(1, 2) # (batch, head, time1, d_k) + k = k.transpose(1, 2) # (batch, head, time2, d_k) + v = v.transpose(1, 2) # (batch, head, time2, d_k) + + return q, k, v + + def forward_attention(self, value, scores, mask): + """Compute attention context vector. + + Args: + value (torch.Tensor): Transformed value (#batch, n_head, time2, d_k). + scores (torch.Tensor): Attention score (#batch, n_head, time1, time2). + mask (torch.Tensor): Mask (#batch, 1, time2) or (#batch, time1, time2). + + Returns: + torch.Tensor: Transformed value (#batch, time1, d_model) + weighted by the attention score (#batch, time1, time2). + + """ + n_batch = value.size(0) + if mask is not None: + mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) + min_value = float(numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min) + scores = scores.masked_fill(mask, min_value) + attn = torch.softmax(scores, dim=-1).masked_fill( + mask, 0.0 + ) # (batch, head, time1, time2) + else: + attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2) + + p_attn = self.dropout(attn) + x = torch.matmul(p_attn, value) # (batch, head, time1, d_k) + x = ( + x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) + ) # (batch, time1, d_model) + + return self.linear_out(x) # (batch, time1, d_model) + + def forward(self, query, key, value, mask): + """Compute scaled dot product attention. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + mask (torch.Tensor): Mask tensor (#batch, 1, time2) or + (#batch, time1, time2). + + Returns: + torch.Tensor: Output tensor (#batch, time1, d_model). + + """ + q, k, v = self.forward_qkv(query, key, value) + scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) + return self.forward_attention(v, scores, mask) + + +class RelPositionMultiHeadedAttention(MultiHeadedAttention): + """Multi-Head Attention layer with relative position encoding (new implementation). + + Details can be found in https://github.com/espnet/espnet/pull/2816. + + Paper: https://arxiv.org/abs/1901.02860 + + Args: + n_head (int): The number of heads. + n_feat (int): The number of features. + dropout_rate (float): Dropout rate. + zero_triu (bool): Whether to zero the upper triangular part of attention matrix. + + """ + + def __init__(self, n_head, n_feat, dropout_rate, zero_triu=False): + """Construct an RelPositionMultiHeadedAttention object.""" + super().__init__(n_head, n_feat, dropout_rate) + self.zero_triu = zero_triu + # linear transformation for positional encoding + self.linear_pos = nn.Linear(n_feat, n_feat, bias=False) + # these two learnable bias are used in matrix c and matrix d + # as described in https://arxiv.org/abs/1901.02860 Section 3.3 + self.pos_bias_u = nn.Parameter(torch.Tensor(self.h, self.d_k)) + self.pos_bias_v = nn.Parameter(torch.Tensor(self.h, self.d_k)) + torch.nn.init.xavier_uniform_(self.pos_bias_u) + torch.nn.init.xavier_uniform_(self.pos_bias_v) + + def rel_shift(self, x): + """Compute relative positional encoding. + + Args: + x (torch.Tensor): Input tensor (batch, head, time1, 2*time1-1). + time1 means the length of query vector. + + Returns: + torch.Tensor: Output tensor. + + """ + zero_pad = torch.zeros((*x.size()[:3], 1), device=x.device, dtype=x.dtype) + x_padded = torch.cat([zero_pad, x], dim=-1) + + x_padded = x_padded.view(*x.size()[:2], x.size(3) + 1, x.size(2)) + x = x_padded[:, :, 1:].view_as(x)[ + :, :, :, : x.size(-1) // 2 + 1 + ] # only keep the positions from 0 to time2 + + if self.zero_triu: + ones = torch.ones((x.size(2), x.size(3)), device=x.device) + x = x * torch.tril(ones, x.size(3) - x.size(2))[None, None, :, :] + + return x + + def forward(self, query, key, value, pos_emb, mask): + """Compute 'Scaled Dot Product Attention' with rel. positional encoding. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + pos_emb (torch.Tensor): Positional embedding tensor + (#batch, 2*time1-1, size). + mask (torch.Tensor): Mask tensor (#batch, 1, time2) or + (#batch, time1, time2). + + Returns: + torch.Tensor: Output tensor (#batch, time1, d_model). + + """ + q, k, v = self.forward_qkv(query, key, value) + q = q.transpose(1, 2) # (batch, time1, head, d_k) + + n_batch_pos = pos_emb.size(0) + p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k) + p = p.transpose(1, 2) # (batch, head, 2*time1-1, d_k) + + # (batch, head, time1, d_k) + q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2) + # (batch, head, time1, d_k) + q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2) + + # compute attention score + # first compute matrix a and matrix c + # as described in https://arxiv.org/abs/1901.02860 Section 3.3 + # (batch, head, time1, time2) + matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1)) + + # compute matrix b and matrix d + # (batch, head, time1, 2*time1-1) + matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) + matrix_bd = self.rel_shift(matrix_bd) + + scores = (matrix_ac + matrix_bd) / math.sqrt(self.d_k) # (batch, head, time1, time2) + + return self.forward_attention(v, scores, mask) + + +class MultiHeadSelfAttention(nn.Module): + """Multi-Head Attention layer. + + Args: + n_head (int): The number of heads. + n_feat (int): The number of features. + dropout_rate (float): Dropout rate. + + """ + + def __init__(self, n_head, in_feat, n_feat, dropout_rate): + """Construct an MultiHeadedAttention object.""" + super(MultiHeadSelfAttention, self).__init__() + assert n_feat % n_head == 0 + # We assume d_v always equals d_k + self.d_k = n_feat // n_head + self.h = n_head + self.linear_out = nn.Linear(n_feat, n_feat) + self.linear_q_k_v = nn.Linear(in_feat, n_feat * 3) + self.attn = None + self.dropout = nn.Dropout(p=dropout_rate) + + def forward_qkv(self, x): + """Transform query, key and value. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + + Returns: + torch.Tensor: Transformed query tensor (#batch, n_head, time1, d_k). + torch.Tensor: Transformed key tensor (#batch, n_head, time2, d_k). + torch.Tensor: Transformed value tensor (#batch, n_head, time2, d_k). + + """ + b, t, d = x.size() + q_k_v = self.linear_q_k_v(x) + q, k, v = torch.split(q_k_v, int(self.h * self.d_k), dim=-1) + q_h = torch.reshape(q, (b, t, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time1, d_k) + k_h = torch.reshape(k, (b, t, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time2, d_k) + v_h = torch.reshape(v, (b, t, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time2, d_k) + + return q_h, k_h, v_h, v + + def forward_attention(self, value, scores, mask, mask_att_chunk_encoder=None): + """Compute attention context vector. + + Args: + value (torch.Tensor): Transformed value (#batch, n_head, time2, d_k). + scores (torch.Tensor): Attention score (#batch, n_head, time1, time2). + mask (torch.Tensor): Mask (#batch, 1, time2) or (#batch, time1, time2). + + Returns: + torch.Tensor: Transformed value (#batch, time1, d_model) + weighted by the attention score (#batch, time1, time2). + + """ + n_batch = value.size(0) + if mask is not None: + if mask_att_chunk_encoder is not None: + mask = mask * mask_att_chunk_encoder + + mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) + + min_value = float(numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min) + scores = scores.masked_fill(mask, min_value) + attn = torch.softmax(scores, dim=-1).masked_fill( + mask, 0.0 + ) # (batch, head, time1, time2) + else: + attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2) + + p_attn = self.dropout(attn) + x = torch.matmul(p_attn, value) # (batch, head, time1, d_k) + x = ( + x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) + ) # (batch, time1, d_model) + + return self.linear_out(x) # (batch, time1, d_model) + + def forward(self, x, mask, mask_att_chunk_encoder=None): + """Compute scaled dot product attention. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + mask (torch.Tensor): Mask tensor (#batch, 1, time2) or + (#batch, time1, time2). + + Returns: + torch.Tensor: Output tensor (#batch, time1, d_model). + + """ + q_h, k_h, v_h, v = self.forward_qkv(x) + q_h = q_h * self.d_k ** (-0.5) + scores = torch.matmul(q_h, k_h.transpose(-2, -1)) + att_outs = self.forward_attention(v_h, scores, mask, mask_att_chunk_encoder) + return att_outs diff --git a/funasr/models/sond/e2e_diar_sond.py b/funasr/models/sond/e2e_diar_sond.py new file mode 100644 index 0000000000000000000000000000000000000000..0f5a7c27c2d4cf4850b7a75ce33c6e354de8904c --- /dev/null +++ b/funasr/models/sond/e2e_diar_sond.py @@ -0,0 +1,576 @@ +#!/usr/bin/env python3 +# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved. +# MIT License (https://opensource.org/licenses/MIT) +import logging +import random +from contextlib import contextmanager +from distutils.version import LooseVersion +from itertools import permutations +from typing import Dict +from typing import Optional +from typing import Tuple, List + +import numpy as np +import torch +from torch.nn import functional as F + +from funasr.models.transformer.utils.nets_utils import to_device +from funasr.models.transformer.utils.nets_utils import make_pad_mask +from funasr.models.decoder.abs_decoder import AbsDecoder +from funasr.models.encoder.abs_encoder import AbsEncoder +from funasr.frontends.abs_frontend import AbsFrontend +from funasr.models.specaug.abs_specaug import AbsSpecAug +from funasr.models.specaug.abs_profileaug import AbsProfileAug +from funasr.layers.abs_normalize import AbsNormalize +from funasr.train_utils.device_funcs import force_gatherable +from funasr.models.base_model import FunASRModel +from funasr.losses.label_smoothing_loss import LabelSmoothingLoss, SequenceBinaryCrossEntropy +from funasr.utils.misc import int2vec +from funasr.utils.hinter import hint_once + +if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): + from torch.cuda.amp import autocast +else: + # Nothing to do if torch<1.6.0 + @contextmanager + def autocast(enabled=True): + yield + + +class DiarSondModel(FunASRModel): + """Speaker overlap-aware neural diarization model + reference: https://arxiv.org/abs/2211.10243 + """ + + def __init__( + self, + vocab_size: int, + frontend: Optional[AbsFrontend], + specaug: Optional[AbsSpecAug], + profileaug: Optional[AbsProfileAug], + normalize: Optional[AbsNormalize], + encoder: torch.nn.Module, + speaker_encoder: Optional[torch.nn.Module], + ci_scorer: torch.nn.Module, + cd_scorer: Optional[torch.nn.Module], + decoder: torch.nn.Module, + token_list: list, + lsm_weight: float = 0.1, + length_normalized_loss: bool = False, + max_spk_num: int = 16, + label_aggregator: Optional[torch.nn.Module] = None, + normalize_speech_speaker: bool = False, + ignore_id: int = -1, + speaker_discrimination_loss_weight: float = 1.0, + inter_score_loss_weight: float = 0.0, + inputs_type: str = "raw", + model_regularizer_weight: float = 0.0, + freeze_encoder: bool = False, + onfly_shuffle_speaker: bool = True, + ): + + super().__init__() + + self.encoder = encoder + self.speaker_encoder = speaker_encoder + self.ci_scorer = ci_scorer + self.cd_scorer = cd_scorer + self.normalize = normalize + self.frontend = frontend + self.specaug = specaug + self.profileaug = profileaug + self.label_aggregator = label_aggregator + self.decoder = decoder + self.token_list = token_list + self.max_spk_num = max_spk_num + self.normalize_speech_speaker = normalize_speech_speaker + self.ignore_id = ignore_id + self.model_regularizer_weight = model_regularizer_weight + self.freeze_encoder = freeze_encoder + self.onfly_shuffle_speaker = onfly_shuffle_speaker + self.criterion_diar = LabelSmoothingLoss( + size=vocab_size, + padding_idx=ignore_id, + smoothing=lsm_weight, + normalize_length=length_normalized_loss, + ) + self.criterion_bce = SequenceBinaryCrossEntropy(normalize_length=length_normalized_loss) + self.pse_embedding = self.generate_pse_embedding() + self.power_weight = torch.from_numpy( + 2 ** np.arange(max_spk_num)[np.newaxis, np.newaxis, :] + ).float() + self.int_token_arr = torch.from_numpy( + np.array(self.token_list).astype(int)[np.newaxis, np.newaxis, :] + ).int() + self.speaker_discrimination_loss_weight = speaker_discrimination_loss_weight + self.inter_score_loss_weight = inter_score_loss_weight + self.forward_steps = 0 + self.inputs_type = inputs_type + self.to_regularize_parameters = None + + def get_regularize_parameters(self): + to_regularize_parameters, normal_parameters = [], [] + for name, param in self.named_parameters(): + if ( + "encoder" in name + and "weight" in name + and "bn" not in name + and ("conv2" in name or "conv1" in name or "conv_sc" in name or "dense" in name) + ): + to_regularize_parameters.append((name, param)) + else: + normal_parameters.append((name, param)) + self.to_regularize_parameters = to_regularize_parameters + return to_regularize_parameters, normal_parameters + + def generate_pse_embedding(self): + embedding = np.zeros((len(self.token_list), self.max_spk_num), dtype=np.float32) + for idx, pse_label in enumerate(self.token_list): + emb = int2vec(int(pse_label), vec_dim=self.max_spk_num, dtype=np.float32) + embedding[idx] = emb + return torch.from_numpy(embedding) + + def rand_permute_speaker(self, raw_profile, raw_binary_labels): + """ + raw_profile: B, N, D + raw_binary_labels: B, T, N + """ + assert ( + raw_profile.shape[1] == raw_binary_labels.shape[2] + ), "Num profile: {}, Num label: {}".format( + raw_profile.shape[1], raw_binary_labels.shape[-1] + ) + profile = torch.clone(raw_profile) + binary_labels = torch.clone(raw_binary_labels) + bsz, num_spk = profile.shape[0], profile.shape[1] + for i in range(bsz): + idx = list(range(num_spk)) + random.shuffle(idx) + profile[i] = profile[i][idx, :] + binary_labels[i] = binary_labels[i][:, idx] + + return profile, binary_labels + + def forward( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor = None, + profile: torch.Tensor = None, + profile_lengths: torch.Tensor = None, + binary_labels: torch.Tensor = None, + binary_labels_lengths: torch.Tensor = None, + ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]: + """Frontend + Encoder + Speaker Encoder + CI Scorer + CD Scorer + Decoder + Calc loss + + Args: + speech: (Batch, samples) or (Batch, frames, input_size) + speech_lengths: (Batch,) default None for chunk interator, + because the chunk-iterator does not + have the speech_lengths returned. + see in + espnet2/iterators/chunk_iter_factory.py + profile: (Batch, N_spk, dim) + profile_lengths: (Batch,) + binary_labels: (Batch, frames, max_spk_num) + binary_labels_lengths: (Batch,) + """ + assert speech.shape[0] <= binary_labels.shape[0], (speech.shape, binary_labels.shape) + batch_size = speech.shape[0] + if self.freeze_encoder: + hint_once("Freeze encoder", "freeze_encoder", rank=0) + self.encoder.eval() + self.forward_steps = self.forward_steps + 1 + if self.pse_embedding.device != speech.device: + self.pse_embedding = self.pse_embedding.to(speech.device) + self.power_weight = self.power_weight.to(speech.device) + self.int_token_arr = self.int_token_arr.to(speech.device) + + if self.onfly_shuffle_speaker: + hint_once("On-the-fly shuffle speaker permutation.", "onfly_shuffle_speaker", rank=0) + profile, binary_labels = self.rand_permute_speaker(profile, binary_labels) + + # 0a. Aggregate time-domain labels to match forward outputs + if self.label_aggregator is not None: + binary_labels, binary_labels_lengths = self.label_aggregator( + binary_labels, binary_labels_lengths + ) + # 0b. augment profiles + if self.profileaug is not None and self.training: + speech, profile, binary_labels = self.profileaug( + speech, + speech_lengths, + profile, + profile_lengths, + binary_labels, + binary_labels_lengths, + ) + + # 1. Calculate power-set encoding (PSE) labels + pad_bin_labels = F.pad( + binary_labels, (0, self.max_spk_num - binary_labels.shape[2]), "constant", 0.0 + ) + raw_pse_labels = torch.sum(pad_bin_labels * self.power_weight, dim=2, keepdim=True) + pse_labels = torch.argmax((raw_pse_labels.int() == self.int_token_arr).float(), dim=2) + + # 2. Network forward + pred, inter_outputs = self.prediction_forward( + speech, speech_lengths, profile, profile_lengths, return_inter_outputs=True + ) + (speech, speech_lengths), (profile, profile_lengths), (ci_score, cd_score) = inter_outputs + + # If encoder uses conv* as input_layer (i.e., subsampling), + # the sequence length of 'pred' might be slightly less than the + # length of 'spk_labels'. Here we force them to be equal. + length_diff_tolerance = 2 + length_diff = abs(pse_labels.shape[1] - pred.shape[1]) + if length_diff <= length_diff_tolerance: + min_len = min(pred.shape[1], pse_labels.shape[1]) + pse_labels = pse_labels[:, :min_len] + pred = pred[:, :min_len] + cd_score = cd_score[:, :min_len] + ci_score = ci_score[:, :min_len] + + loss_diar = self.classification_loss(pred, pse_labels, binary_labels_lengths) + loss_spk_dis = self.speaker_discrimination_loss(profile, profile_lengths) + loss_inter_ci, loss_inter_cd = self.internal_score_loss( + cd_score, ci_score, pse_labels, binary_labels_lengths + ) + regularizer_loss = None + if self.model_regularizer_weight > 0 and self.to_regularize_parameters is not None: + regularizer_loss = self.calculate_regularizer_loss() + label_mask = make_pad_mask(binary_labels_lengths, maxlen=pse_labels.shape[1]).to( + pse_labels.device + ) + loss = ( + loss_diar + + self.speaker_discrimination_loss_weight * loss_spk_dis + + self.inter_score_loss_weight * (loss_inter_ci + loss_inter_cd) + ) + # if regularizer_loss is not None: + # loss = loss + regularizer_loss * self.model_regularizer_weight + + ( + correct, + num_frames, + speech_scored, + speech_miss, + speech_falarm, + speaker_scored, + speaker_miss, + speaker_falarm, + speaker_error, + ) = self.calc_diarization_error( + pred=F.embedding(pred.argmax(dim=2) * (~label_mask), self.pse_embedding), + label=F.embedding(pse_labels * (~label_mask), self.pse_embedding), + length=binary_labels_lengths, + ) + + if speech_scored > 0 and num_frames > 0: + sad_mr, sad_fr, mi, fa, cf, acc, der = ( + speech_miss / speech_scored, + speech_falarm / speech_scored, + speaker_miss / speaker_scored, + speaker_falarm / speaker_scored, + speaker_error / speaker_scored, + correct / num_frames, + (speaker_miss + speaker_falarm + speaker_error) / speaker_scored, + ) + else: + sad_mr, sad_fr, mi, fa, cf, acc, der = 0, 0, 0, 0, 0, 0, 0 + + stats = dict( + loss=loss.detach(), + loss_diar=loss_diar.detach() if loss_diar is not None else None, + loss_spk_dis=loss_spk_dis.detach() if loss_spk_dis is not None else None, + loss_inter_ci=loss_inter_ci.detach() if loss_inter_ci is not None else None, + loss_inter_cd=loss_inter_cd.detach() if loss_inter_cd is not None else None, + regularizer_loss=regularizer_loss.detach() if regularizer_loss is not None else None, + sad_mr=sad_mr, + sad_fr=sad_fr, + mi=mi, + fa=fa, + cf=cf, + acc=acc, + der=der, + forward_steps=self.forward_steps, + ) + + loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) + return loss, stats, weight + + def calculate_regularizer_loss(self): + regularizer_loss = 0.0 + for name, param in self.to_regularize_parameters: + regularizer_loss = regularizer_loss + torch.norm(param, p=2) + return regularizer_loss + + def classification_loss( + self, predictions: torch.Tensor, labels: torch.Tensor, prediction_lengths: torch.Tensor + ) -> torch.Tensor: + mask = make_pad_mask(prediction_lengths, maxlen=labels.shape[1]) + pad_labels = labels.masked_fill(mask.to(predictions.device), value=self.ignore_id) + loss = self.criterion_diar(predictions.contiguous(), pad_labels) + + return loss + + def speaker_discrimination_loss( + self, profile: torch.Tensor, profile_lengths: torch.Tensor + ) -> torch.Tensor: + profile_mask = ( + torch.linalg.norm(profile, ord=2, dim=2, keepdim=True) > 0 + ).float() # (B, N, 1) + mask = torch.matmul(profile_mask, profile_mask.transpose(1, 2)) # (B, N, N) + mask = mask * (1.0 - torch.eye(self.max_spk_num).unsqueeze(0).to(mask)) + + eps = 1e-12 + coding_norm = ( + torch.linalg.norm( + profile * profile_mask + (1 - profile_mask) * eps, dim=2, keepdim=True + ) + * profile_mask + ) + # profile: Batch, N, dim + cos_theta = ( + F.cosine_similarity(profile.unsqueeze(2), profile.unsqueeze(1), dim=-1, eps=eps) * mask + ) + cos_theta = torch.clip(cos_theta, -1 + eps, 1 - eps) + loss = (F.relu(mask * coding_norm * (cos_theta - 0.0))).sum() / mask.sum() + + return loss + + def calculate_multi_labels(self, pse_labels, pse_labels_lengths): + mask = make_pad_mask(pse_labels_lengths, maxlen=pse_labels.shape[1]) + padding_labels = pse_labels.masked_fill(mask.to(pse_labels.device), value=0).to(pse_labels) + multi_labels = F.embedding(padding_labels, self.pse_embedding) + + return multi_labels + + def internal_score_loss( + self, + cd_score: torch.Tensor, + ci_score: torch.Tensor, + pse_labels: torch.Tensor, + pse_labels_lengths: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + multi_labels = self.calculate_multi_labels(pse_labels, pse_labels_lengths) + ci_loss = self.criterion_bce(ci_score, multi_labels, pse_labels_lengths) + cd_loss = self.criterion_bce(cd_score, multi_labels, pse_labels_lengths) + return ci_loss, cd_loss + + def collect_feats( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + profile: torch.Tensor = None, + profile_lengths: torch.Tensor = None, + binary_labels: torch.Tensor = None, + binary_labels_lengths: torch.Tensor = None, + ) -> Dict[str, torch.Tensor]: + feats, feats_lengths = self._extract_feats(speech, speech_lengths) + return {"feats": feats, "feats_lengths": feats_lengths} + + def encode_speaker( + self, + profile: torch.Tensor, + profile_lengths: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + with autocast(False): + if profile.shape[1] < self.max_spk_num: + profile = F.pad( + profile, [0, 0, 0, self.max_spk_num - profile.shape[1], 0, 0], "constant", 0.0 + ) + profile_mask = (torch.linalg.norm(profile, ord=2, dim=2, keepdim=True) > 0).float() + profile = F.normalize(profile, dim=2) + if self.speaker_encoder is not None: + profile = self.speaker_encoder(profile, profile_lengths)[0] + return profile * profile_mask, profile_lengths + else: + return profile, profile_lengths + + def encode_speech( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + if self.encoder is not None and self.inputs_type == "raw": + speech, speech_lengths = self.encode(speech, speech_lengths) + speech_mask = ~make_pad_mask(speech_lengths, maxlen=speech.shape[1]) + speech_mask = speech_mask.to(speech.device).unsqueeze(-1).float() + return speech * speech_mask, speech_lengths + else: + return speech, speech_lengths + + @staticmethod + def concate_speech_ivc(speech: torch.Tensor, ivc: torch.Tensor) -> torch.Tensor: + nn, tt = ivc.shape[1], speech.shape[1] + speech = speech.unsqueeze(dim=1) # B x 1 x T x D + speech = speech.expand(-1, nn, -1, -1) # B x N x T x D + ivc = ivc.unsqueeze(dim=2) # B x N x 1 x D + ivc = ivc.expand(-1, -1, tt, -1) # B x N x T x D + sd_in = torch.cat([speech, ivc], dim=3) # B x N x T x 2D + return sd_in + + def calc_similarity( + self, + speech_encoder_outputs: torch.Tensor, + speaker_encoder_outputs: torch.Tensor, + seq_len: torch.Tensor = None, + spk_len: torch.Tensor = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + bb, tt = speech_encoder_outputs.shape[0], speech_encoder_outputs.shape[1] + d_sph, d_spk = speech_encoder_outputs.shape[2], speaker_encoder_outputs.shape[2] + if self.normalize_speech_speaker: + speech_encoder_outputs = F.normalize(speech_encoder_outputs, dim=2) + speaker_encoder_outputs = F.normalize(speaker_encoder_outputs, dim=2) + ge_in = self.concate_speech_ivc(speech_encoder_outputs, speaker_encoder_outputs) + ge_in = torch.reshape(ge_in, [bb * self.max_spk_num, tt, d_sph + d_spk]) + ge_len = seq_len.unsqueeze(1).expand(-1, self.max_spk_num) + ge_len = torch.reshape(ge_len, [bb * self.max_spk_num]) + cd_simi = self.cd_scorer(ge_in, ge_len)[0] + cd_simi = torch.reshape(cd_simi, [bb, self.max_spk_num, tt, 1]) + cd_simi = cd_simi.squeeze(dim=3).permute([0, 2, 1]) + + if isinstance(self.ci_scorer, AbsEncoder): + ci_simi = self.ci_scorer(ge_in, ge_len)[0] + ci_simi = torch.reshape(ci_simi, [bb, self.max_spk_num, tt]).permute([0, 2, 1]) + else: + ci_simi = self.ci_scorer(speech_encoder_outputs, speaker_encoder_outputs) + + return ci_simi, cd_simi + + def post_net_forward(self, simi, seq_len): + logits = self.decoder(simi, seq_len)[0] + + return logits + + def prediction_forward( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + profile: torch.Tensor, + profile_lengths: torch.Tensor, + return_inter_outputs: bool = False, + ) -> [torch.Tensor, Optional[list]]: + # speech encoding + speech, speech_lengths = self.encode_speech(speech, speech_lengths) + # speaker encoding + profile, profile_lengths = self.encode_speaker(profile, profile_lengths) + # calculating similarity + ci_simi, cd_simi = self.calc_similarity(speech, profile, speech_lengths, profile_lengths) + similarity = torch.cat([cd_simi, ci_simi], dim=2) + # post net forward + logits = self.post_net_forward(similarity, speech_lengths) + + if return_inter_outputs: + return logits, [ + (speech, speech_lengths), + (profile, profile_lengths), + (ci_simi, cd_simi), + ] + return logits + + def encode( + self, speech: torch.Tensor, speech_lengths: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Frontend + Encoder + + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch,) + """ + with autocast(False): + # 1. Extract feats + feats, feats_lengths = self._extract_feats(speech, speech_lengths) + + # 2. Data augmentation + if self.specaug is not None and self.training: + feats, feats_lengths = self.specaug(feats, feats_lengths) + + # 3. Normalization for feature: e.g. Global-CMVN, Utterance-CMVN + if self.normalize is not None: + feats, feats_lengths = self.normalize(feats, feats_lengths) + + # 4. Forward encoder + # feats: (Batch, Length, Dim) + # -> encoder_out: (Batch, Length2, Dim) + encoder_outputs = self.encoder(feats, feats_lengths) + encoder_out, encoder_out_lens = encoder_outputs[:2] + + assert encoder_out.size(0) == speech.size(0), ( + encoder_out.size(), + speech.size(0), + ) + assert encoder_out.size(1) <= encoder_out_lens.max(), ( + encoder_out.size(), + encoder_out_lens.max(), + ) + + return encoder_out, encoder_out_lens + + def _extract_feats( + self, speech: torch.Tensor, speech_lengths: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + batch_size = speech.shape[0] + speech_lengths = ( + speech_lengths + if speech_lengths is not None + else torch.ones(batch_size).int() * speech.shape[1] + ) + + assert speech_lengths.dim() == 1, speech_lengths.shape + + # for data-parallel + speech = speech[:, : speech_lengths.max()] + + if self.frontend is not None: + # Frontend + # e.g. STFT and Feature extract + # data_loader may send time-domain signal in this case + # speech (Batch, NSamples) -> feats: (Batch, NFrames, Dim) + feats, feats_lengths = self.frontend(speech, speech_lengths) + else: + # No frontend and no feature extract + feats, feats_lengths = speech, speech_lengths + return feats, feats_lengths + + @staticmethod + def calc_diarization_error(pred, label, length): + # Note (jiatong): Credit to https://github.com/hitachi-speech/EEND + + (batch_size, max_len, num_output) = label.size() + # mask the padding part + mask = ~make_pad_mask(length, maxlen=label.shape[1]).unsqueeze(-1).numpy() + + # pred and label have the shape (batch_size, max_len, num_output) + label_np = label.data.cpu().numpy().astype(int) + pred_np = (pred.data.cpu().numpy() > 0).astype(int) + label_np = label_np * mask + pred_np = pred_np * mask + length = length.data.cpu().numpy() + + # compute speech activity detection error + n_ref = np.sum(label_np, axis=2) + n_sys = np.sum(pred_np, axis=2) + speech_scored = float(np.sum(n_ref > 0)) + speech_miss = float(np.sum(np.logical_and(n_ref > 0, n_sys == 0))) + speech_falarm = float(np.sum(np.logical_and(n_ref == 0, n_sys > 0))) + + # compute speaker diarization error + speaker_scored = float(np.sum(n_ref)) + speaker_miss = float(np.sum(np.maximum(n_ref - n_sys, 0))) + speaker_falarm = float(np.sum(np.maximum(n_sys - n_ref, 0))) + n_map = np.sum(np.logical_and(label_np == 1, pred_np == 1), axis=2) + speaker_error = float(np.sum(np.minimum(n_ref, n_sys) - n_map)) + correct = float(1.0 * np.sum((label_np == pred_np) * mask) / num_output) + num_frames = np.sum(length) + return ( + correct, + num_frames, + speech_scored, + speech_miss, + speech_falarm, + speaker_scored, + speaker_miss, + speaker_falarm, + speaker_error, + ) diff --git a/funasr/models/sond/encoder/__init__.py b/funasr/models/sond/encoder/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/models/sond/encoder/ci_scorers.py b/funasr/models/sond/encoder/ci_scorers.py new file mode 100644 index 0000000000000000000000000000000000000000..a60560b7c1b445b49261a9d194773e5223ace08d --- /dev/null +++ b/funasr/models/sond/encoder/ci_scorers.py @@ -0,0 +1,32 @@ +import torch +from torch.nn import functional as F + + +class DotScorer(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward( + self, + xs_pad: torch.Tensor, + spk_emb: torch.Tensor, + ): + # xs_pad: B, T, D + # spk_emb: B, N, D + scores = torch.matmul(xs_pad, spk_emb.transpose(1, 2)) + return scores + + +class CosScorer(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward( + self, + xs_pad: torch.Tensor, + spk_emb: torch.Tensor, + ): + # xs_pad: B, T, D + # spk_emb: B, N, D + scores = F.cosine_similarity(xs_pad.unsqueeze(2), spk_emb.unsqueeze(1), dim=-1) + return scores diff --git a/funasr/models/sond/encoder/conv_encoder.py b/funasr/models/sond/encoder/conv_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..c4f7098c7f227676c31d33a5b4f0dfdbb12dbae3 --- /dev/null +++ b/funasr/models/sond/encoder/conv_encoder.py @@ -0,0 +1,174 @@ +from typing import List +from typing import Optional +from typing import Sequence +from typing import Tuple +from typing import Union +import logging +import torch +import torch.nn as nn +from torch.nn import functional as F +import numpy as np +from funasr.models.transformer.utils.nets_utils import make_pad_mask +from funasr.models.transformer.layer_norm import LayerNorm +from funasr.models.encoder.abs_encoder import AbsEncoder +import math +from funasr.models.transformer.utils.repeat import repeat + + +class EncoderLayer(nn.Module): + def __init__( + self, + input_units, + num_units, + kernel_size=3, + activation="tanh", + stride=1, + include_batch_norm=False, + residual=False, + ): + super().__init__() + left_padding = math.ceil((kernel_size - stride) / 2) + right_padding = kernel_size - stride - left_padding + self.conv_padding = nn.ConstantPad1d((left_padding, right_padding), 0.0) + self.conv1d = nn.Conv1d( + input_units, + num_units, + kernel_size, + stride, + ) + self.activation = self.get_activation(activation) + if include_batch_norm: + self.bn = nn.BatchNorm1d(num_units, momentum=0.99, eps=1e-3) + self.residual = residual + self.include_batch_norm = include_batch_norm + self.input_units = input_units + self.num_units = num_units + self.stride = stride + + @staticmethod + def get_activation(activation): + if activation == "tanh": + return nn.Tanh() + else: + return nn.ReLU() + + def forward(self, xs_pad, ilens=None): + outputs = self.conv1d(self.conv_padding(xs_pad)) + if self.residual and self.stride == 1 and self.input_units == self.num_units: + outputs = outputs + xs_pad + + if self.include_batch_norm: + outputs = self.bn(outputs) + + # add parenthesis for repeat module + return self.activation(outputs), ilens + + +class ConvEncoder(AbsEncoder): + """ + Author: Speech Lab of DAMO Academy, Alibaba Group + Convolution encoder in OpenNMT framework + """ + + def __init__( + self, + num_layers, + input_units, + num_units, + kernel_size=3, + dropout_rate=0.3, + position_encoder=None, + activation="tanh", + auxiliary_states=True, + out_units=None, + out_norm=False, + out_residual=False, + include_batchnorm=False, + regularization_weight=0.0, + stride=1, + tf2torch_tensor_name_prefix_torch: str = "speaker_encoder", + tf2torch_tensor_name_prefix_tf: str = "EAND/speaker_encoder", + ): + super().__init__() + self._output_size = num_units + + self.num_layers = num_layers + self.input_units = input_units + self.num_units = num_units + self.kernel_size = kernel_size + self.dropout_rate = dropout_rate + self.position_encoder = position_encoder + self.out_units = out_units + self.auxiliary_states = auxiliary_states + self.out_norm = out_norm + self.activation = activation + self.out_residual = out_residual + self.include_batch_norm = include_batchnorm + self.regularization_weight = regularization_weight + self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch + self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf + if isinstance(stride, int): + self.stride = [stride] * self.num_layers + else: + self.stride = stride + self.downsample_rate = 1 + for s in self.stride: + self.downsample_rate *= s + + self.dropout = nn.Dropout(dropout_rate) + self.cnn_a = repeat( + self.num_layers, + lambda lnum: EncoderLayer( + input_units if lnum == 0 else num_units, + num_units, + kernel_size, + activation, + self.stride[lnum], + include_batchnorm, + residual=True if lnum > 0 else False, + ), + ) + + if self.out_units is not None: + left_padding = math.ceil((kernel_size - stride) / 2) + right_padding = kernel_size - stride - left_padding + self.out_padding = nn.ConstantPad1d((left_padding, right_padding), 0.0) + self.conv_out = nn.Conv1d( + num_units, + out_units, + kernel_size, + ) + + if self.out_norm: + self.after_norm = LayerNorm(out_units) + + def output_size(self) -> int: + return self.num_units + + def forward( + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + prev_states: torch.Tensor = None, + ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + + inputs = xs_pad + if self.position_encoder is not None: + inputs = self.position_encoder(inputs) + + if self.dropout_rate > 0: + inputs = self.dropout(inputs) + + outputs, _ = self.cnn_a(inputs.transpose(1, 2), ilens) + + if self.out_units is not None: + outputs = self.conv_out(self.out_padding(outputs)) + + outputs = outputs.transpose(1, 2) + if self.out_norm: + outputs = self.after_norm(outputs) + + if self.out_residual: + outputs = outputs + inputs + + return outputs, ilens, None diff --git a/funasr/models/sond/encoder/ecapa_tdnn_encoder.py b/funasr/models/sond/encoder/ecapa_tdnn_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..1af8b7033268543c3fb9ee8a5fa4a75b5a33289c --- /dev/null +++ b/funasr/models/sond/encoder/ecapa_tdnn_encoder.py @@ -0,0 +1,672 @@ +import math +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class _BatchNorm1d(nn.Module): + def __init__( + self, + input_shape=None, + input_size=None, + eps=1e-05, + momentum=0.1, + affine=True, + track_running_stats=True, + combine_batch_time=False, + skip_transpose=False, + ): + super().__init__() + self.combine_batch_time = combine_batch_time + self.skip_transpose = skip_transpose + + if input_size is None and skip_transpose: + input_size = input_shape[1] + elif input_size is None: + input_size = input_shape[-1] + + self.norm = nn.BatchNorm1d( + input_size, + eps=eps, + momentum=momentum, + affine=affine, + track_running_stats=track_running_stats, + ) + + def forward(self, x): + shape_or = x.shape + if self.combine_batch_time: + if x.ndim == 3: + x = x.reshape(shape_or[0] * shape_or[1], shape_or[2]) + else: + x = x.reshape(shape_or[0] * shape_or[1], shape_or[3], shape_or[2]) + + elif not self.skip_transpose: + x = x.transpose(-1, 1) + + x_n = self.norm(x) + + if self.combine_batch_time: + x_n = x_n.reshape(shape_or) + elif not self.skip_transpose: + x_n = x_n.transpose(1, -1) + + return x_n + + +class _Conv1d(nn.Module): + def __init__( + self, + out_channels, + kernel_size, + input_shape=None, + in_channels=None, + stride=1, + dilation=1, + padding="same", + groups=1, + bias=True, + padding_mode="reflect", + skip_transpose=False, + ): + super().__init__() + self.kernel_size = kernel_size + self.stride = stride + self.dilation = dilation + self.padding = padding + self.padding_mode = padding_mode + self.unsqueeze = False + self.skip_transpose = skip_transpose + + if input_shape is None and in_channels is None: + raise ValueError("Must provide one of input_shape or in_channels") + + if in_channels is None: + in_channels = self._check_input_shape(input_shape) + + self.conv = nn.Conv1d( + in_channels, + out_channels, + self.kernel_size, + stride=self.stride, + dilation=self.dilation, + padding=0, + groups=groups, + bias=bias, + ) + + def forward(self, x): + if not self.skip_transpose: + x = x.transpose(1, -1) + + if self.unsqueeze: + x = x.unsqueeze(1) + + if self.padding == "same": + x = self._manage_padding(x, self.kernel_size, self.dilation, self.stride) + + elif self.padding == "causal": + num_pad = (self.kernel_size - 1) * self.dilation + x = F.pad(x, (num_pad, 0)) + + elif self.padding == "valid": + pass + + else: + raise ValueError("Padding must be 'same', 'valid' or 'causal'. Got " + self.padding) + + wx = self.conv(x) + + if self.unsqueeze: + wx = wx.squeeze(1) + + if not self.skip_transpose: + wx = wx.transpose(1, -1) + + return wx + + def _manage_padding( + self, + x, + kernel_size: int, + dilation: int, + stride: int, + ): + # Detecting input shape + L_in = x.shape[-1] + + # Time padding + padding = get_padding_elem(L_in, stride, kernel_size, dilation) + + # Applying padding + x = F.pad(x, padding, mode=self.padding_mode) + + return x + + def _check_input_shape(self, shape): + """Checks the input shape and returns the number of input channels.""" + + if len(shape) == 2: + self.unsqueeze = True + in_channels = 1 + elif self.skip_transpose: + in_channels = shape[1] + elif len(shape) == 3: + in_channels = shape[2] + else: + raise ValueError("conv1d expects 2d, 3d inputs. Got " + str(len(shape))) + + # Kernel size must be odd + if self.kernel_size % 2 == 0: + raise ValueError( + "The field kernel size must be an odd number. Got %s." % (self.kernel_size) + ) + return in_channels + + +def get_padding_elem(L_in: int, stride: int, kernel_size: int, dilation: int): + if stride > 1: + n_steps = math.ceil(((L_in - kernel_size * dilation) / stride) + 1) + L_out = stride * (n_steps - 1) + kernel_size * dilation + padding = [kernel_size // 2, kernel_size // 2] + + else: + L_out = (L_in - dilation * (kernel_size - 1) - 1) // stride + 1 + + padding = [(L_in - L_out) // 2, (L_in - L_out) // 2] + return padding + + +# Skip transpose as much as possible for efficiency +class Conv1d(_Conv1d): + def __init__(self, *args, **kwargs): + super().__init__(skip_transpose=True, *args, **kwargs) + + +class BatchNorm1d(_BatchNorm1d): + def __init__(self, *args, **kwargs): + super().__init__(skip_transpose=True, *args, **kwargs) + + +def length_to_mask(length, max_len=None, dtype=None, device=None): + assert len(length.shape) == 1 + + if max_len is None: + max_len = length.max().long().item() # using arange to generate mask + mask = torch.arange(max_len, device=length.device, dtype=length.dtype).expand( + len(length), max_len + ) < length.unsqueeze(1) + + if dtype is None: + dtype = length.dtype + + if device is None: + device = length.device + + mask = torch.as_tensor(mask, dtype=dtype, device=device) + return mask + + +class TDNNBlock(nn.Module): + def __init__( + self, + in_channels, + out_channels, + kernel_size, + dilation, + activation=nn.ReLU, + groups=1, + ): + super(TDNNBlock, self).__init__() + self.conv = Conv1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + dilation=dilation, + groups=groups, + ) + self.activation = activation() + self.norm = BatchNorm1d(input_size=out_channels) + + def forward(self, x): + return self.norm(self.activation(self.conv(x))) + + +class Res2NetBlock(torch.nn.Module): + """An implementation of Res2NetBlock w/ dilation. + + Arguments + --------- + in_channels : int + The number of channels expected in the input. + out_channels : int + The number of output channels. + scale : int + The scale of the Res2Net block. + kernel_size: int + The kernel size of the Res2Net block. + dilation : int + The dilation of the Res2Net block. + + Example + ------- + >>> inp_tensor = torch.rand([8, 120, 64]).transpose(1, 2) + >>> layer = Res2NetBlock(64, 64, scale=4, dilation=3) + >>> out_tensor = layer(inp_tensor).transpose(1, 2) + >>> out_tensor.shape + torch.Size([8, 120, 64]) + """ + + def __init__(self, in_channels, out_channels, scale=8, kernel_size=3, dilation=1): + super(Res2NetBlock, self).__init__() + assert in_channels % scale == 0 + assert out_channels % scale == 0 + + in_channel = in_channels // scale + hidden_channel = out_channels // scale + + self.blocks = nn.ModuleList( + [ + TDNNBlock( + in_channel, + hidden_channel, + kernel_size=kernel_size, + dilation=dilation, + ) + for i in range(scale - 1) + ] + ) + self.scale = scale + + def forward(self, x): + y = [] + for i, x_i in enumerate(torch.chunk(x, self.scale, dim=1)): + if i == 0: + y_i = x_i + elif i == 1: + y_i = self.blocks[i - 1](x_i) + else: + y_i = self.blocks[i - 1](x_i + y_i) + y.append(y_i) + y = torch.cat(y, dim=1) + return y + + +class SEBlock(nn.Module): + """An implementation of squeeze-and-excitation block. + + Arguments + --------- + in_channels : int + The number of input channels. + se_channels : int + The number of output channels after squeeze. + out_channels : int + The number of output channels. + + Example + ------- + >>> inp_tensor = torch.rand([8, 120, 64]).transpose(1, 2) + >>> se_layer = SEBlock(64, 16, 64) + >>> lengths = torch.rand((8,)) + >>> out_tensor = se_layer(inp_tensor, lengths).transpose(1, 2) + >>> out_tensor.shape + torch.Size([8, 120, 64]) + """ + + def __init__(self, in_channels, se_channels, out_channels): + super(SEBlock, self).__init__() + + self.conv1 = Conv1d(in_channels=in_channels, out_channels=se_channels, kernel_size=1) + self.relu = torch.nn.ReLU(inplace=True) + self.conv2 = Conv1d(in_channels=se_channels, out_channels=out_channels, kernel_size=1) + self.sigmoid = torch.nn.Sigmoid() + + def forward(self, x, lengths=None): + L = x.shape[-1] + if lengths is not None: + mask = length_to_mask(lengths * L, max_len=L, device=x.device) + mask = mask.unsqueeze(1) + total = mask.sum(dim=2, keepdim=True) + s = (x * mask).sum(dim=2, keepdim=True) / total + else: + s = x.mean(dim=2, keepdim=True) + + s = self.relu(self.conv1(s)) + s = self.sigmoid(self.conv2(s)) + + return s * x + + +class AttentiveStatisticsPooling(nn.Module): + """This class implements an attentive statistic pooling layer for each channel. + It returns the concatenated mean and std of the input tensor. + + Arguments + --------- + channels: int + The number of input channels. + attention_channels: int + The number of attention channels. + + Example + ------- + >>> inp_tensor = torch.rand([8, 120, 64]).transpose(1, 2) + >>> asp_layer = AttentiveStatisticsPooling(64) + >>> lengths = torch.rand((8,)) + >>> out_tensor = asp_layer(inp_tensor, lengths).transpose(1, 2) + >>> out_tensor.shape + torch.Size([8, 1, 128]) + """ + + def __init__(self, channels, attention_channels=128, global_context=True): + super().__init__() + + self.eps = 1e-12 + self.global_context = global_context + if global_context: + self.tdnn = TDNNBlock(channels * 3, attention_channels, 1, 1) + else: + self.tdnn = TDNNBlock(channels, attention_channels, 1, 1) + self.tanh = nn.Tanh() + self.conv = Conv1d(in_channels=attention_channels, out_channels=channels, kernel_size=1) + + def forward(self, x, lengths=None): + """Calculates mean and std for a batch (input tensor). + + Arguments + --------- + x : torch.Tensor + Tensor of shape [N, C, L]. + """ + L = x.shape[-1] + + def _compute_statistics(x, m, dim=2, eps=self.eps): + mean = (m * x).sum(dim) + std = torch.sqrt((m * (x - mean.unsqueeze(dim)).pow(2)).sum(dim).clamp(eps)) + return mean, std + + if lengths is None: + lengths = torch.ones(x.shape[0], device=x.device) + + # Make binary mask of shape [N, 1, L] + mask = length_to_mask(lengths * L, max_len=L, device=x.device) + mask = mask.unsqueeze(1) + + # Expand the temporal context of the pooling layer by allowing the + # self-attention to look at global properties of the utterance. + if self.global_context: + # torch.std is unstable for backward computation + # https://github.com/pytorch/pytorch/issues/4320 + total = mask.sum(dim=2, keepdim=True).float() + mean, std = _compute_statistics(x, mask / total) + mean = mean.unsqueeze(2).repeat(1, 1, L) + std = std.unsqueeze(2).repeat(1, 1, L) + attn = torch.cat([x, mean, std], dim=1) + else: + attn = x + + # Apply layers + attn = self.conv(self.tanh(self.tdnn(attn))) + + # Filter out zero-paddings + attn = attn.masked_fill(mask == 0, float("-inf")) + + attn = F.softmax(attn, dim=2) + mean, std = _compute_statistics(x, attn) + # Append mean and std of the batch + pooled_stats = torch.cat((mean, std), dim=1) + pooled_stats = pooled_stats.unsqueeze(2) + + return pooled_stats + + +class SERes2NetBlock(nn.Module): + """An implementation of building block in ECAPA-TDNN, i.e., + TDNN-Res2Net-TDNN-SEBlock. + + Arguments + ---------- + out_channels: int + The number of output channels. + res2net_scale: int + The scale of the Res2Net block. + kernel_size: int + The kernel size of the TDNN blocks. + dilation: int + The dilation of the Res2Net block. + activation : torch class + A class for constructing the activation layers. + groups: int + Number of blocked connections from input channels to output channels. + + Example + ------- + >>> x = torch.rand(8, 120, 64).transpose(1, 2) + >>> conv = SERes2NetBlock(64, 64, res2net_scale=4) + >>> out = conv(x).transpose(1, 2) + >>> out.shape + torch.Size([8, 120, 64]) + """ + + def __init__( + self, + in_channels, + out_channels, + res2net_scale=8, + se_channels=128, + kernel_size=1, + dilation=1, + activation=torch.nn.ReLU, + groups=1, + ): + super().__init__() + self.out_channels = out_channels + self.tdnn1 = TDNNBlock( + in_channels, + out_channels, + kernel_size=1, + dilation=1, + activation=activation, + groups=groups, + ) + self.res2net_block = Res2NetBlock( + out_channels, out_channels, res2net_scale, kernel_size, dilation + ) + self.tdnn2 = TDNNBlock( + out_channels, + out_channels, + kernel_size=1, + dilation=1, + activation=activation, + groups=groups, + ) + self.se_block = SEBlock(out_channels, se_channels, out_channels) + + self.shortcut = None + if in_channels != out_channels: + self.shortcut = Conv1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=1, + ) + + def forward(self, x, lengths=None): + residual = x + if self.shortcut: + residual = self.shortcut(x) + + x = self.tdnn1(x) + x = self.res2net_block(x) + x = self.tdnn2(x) + x = self.se_block(x, lengths) + + return x + residual + + +class ECAPA_TDNN(torch.nn.Module): + """An implementation of the speaker embedding model in a paper. + "ECAPA-TDNN: Emphasized Channel Attention, Propagation and Aggregation in + TDNN Based Speaker Verification" (https://arxiv.org/abs/2005.07143). + + Arguments + --------- + activation : torch class + A class for constructing the activation layers. + channels : list of ints + Output channels for TDNN/SERes2Net layer. + kernel_sizes : list of ints + List of kernel sizes for each layer. + dilations : list of ints + List of dilations for kernels in each layer. + lin_neurons : int + Number of neurons in linear layers. + groups : list of ints + List of groups for kernels in each layer. + + Example + ------- + >>> input_feats = torch.rand([5, 120, 80]) + >>> compute_embedding = ECAPA_TDNN(80, lin_neurons=192) + >>> outputs = compute_embedding(input_feats) + >>> outputs.shape + torch.Size([5, 1, 192]) + """ + + def __init__( + self, + input_size, + lin_neurons=192, + activation=torch.nn.ReLU, + channels=[512, 512, 512, 512, 1536], + kernel_sizes=[5, 3, 3, 3, 1], + dilations=[1, 2, 3, 4, 1], + attention_channels=128, + res2net_scale=8, + se_channels=128, + global_context=True, + groups=[1, 1, 1, 1, 1], + window_size=20, + window_shift=1, + ): + + super().__init__() + assert len(channels) == len(kernel_sizes) + assert len(channels) == len(dilations) + self.channels = channels + self.blocks = nn.ModuleList() + self.window_size = window_size + self.window_shift = window_shift + + # The initial TDNN layer + self.blocks.append( + TDNNBlock( + input_size, + channels[0], + kernel_sizes[0], + dilations[0], + activation, + groups[0], + ) + ) + + # SE-Res2Net layers + for i in range(1, len(channels) - 1): + self.blocks.append( + SERes2NetBlock( + channels[i - 1], + channels[i], + res2net_scale=res2net_scale, + se_channels=se_channels, + kernel_size=kernel_sizes[i], + dilation=dilations[i], + activation=activation, + groups=groups[i], + ) + ) + + # Multi-layer feature aggregation + self.mfa = TDNNBlock( + channels[-1], + channels[-1], + kernel_sizes[-1], + dilations[-1], + activation, + groups=groups[-1], + ) + + # Attentive Statistical Pooling + self.asp = AttentiveStatisticsPooling( + channels[-1], + attention_channels=attention_channels, + global_context=global_context, + ) + self.asp_bn = BatchNorm1d(input_size=channels[-1] * 2) + + # Final linear transformation + self.fc = Conv1d( + in_channels=channels[-1] * 2, + out_channels=lin_neurons, + kernel_size=1, + ) + + def windowed_pooling(self, x, lengths=None): + # x: Batch, Channel, Time + tt = x.shape[2] + num_chunk = int(math.ceil(tt / self.window_shift)) + pad = self.window_size // 2 + x = F.pad(x, (pad, pad, 0, 0), "reflect") + stat_list = [] + + for i in range(num_chunk): + # B x C + st, ed = i * self.window_shift, i * self.window_shift + self.window_size + x = self.asp( + x[:, :, st:ed], + lengths=( + torch.clamp(lengths - i, 0, self.window_size) if lengths is not None else None + ), + ) + x = self.asp_bn(x) + x = self.fc(x) + stat_list.append(x) + + return torch.cat(stat_list, dim=2) + + def forward(self, x, lengths=None): + """Returns the embedding vector. + + Arguments + --------- + x : torch.Tensor + Tensor of shape (batch, time, channel). + lengths: torch.Tensor + Tensor of shape (batch, ) + """ + # Minimize transpose for efficiency + x = x.transpose(1, 2) + + xl = [] + for layer in self.blocks: + try: + x = layer(x, lengths=lengths) + except TypeError: + x = layer(x) + xl.append(x) + + # Multi-layer feature aggregation + x = torch.cat(xl[1:], dim=1) + x = self.mfa(x) + + if self.window_size is None: + # Attentive Statistical Pooling + x = self.asp(x, lengths=lengths) + x = self.asp_bn(x) + # Final linear transformation + x = self.fc(x) + # x = x.transpose(1, 2) + x = x.squeeze(2) # -> B, C + else: + x = self.windowed_pooling(x, lengths) + x = x.transpose(1, 2) # -> B, T, C + return x diff --git a/funasr/models/sond/encoder/fsmn_encoder.py b/funasr/models/sond/encoder/fsmn_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..9ec9912a46116ed013230deda83f8821f822fc91 --- /dev/null +++ b/funasr/models/sond/encoder/fsmn_encoder.py @@ -0,0 +1,180 @@ +from typing import List +from typing import Optional +from typing import Sequence +from typing import Tuple +from typing import Union +import logging +import torch +import torch.nn as nn +from torch.nn import functional as F +import numpy as np +from funasr.models.transformer.utils.nets_utils import make_pad_mask +from funasr.models.transformer.layer_norm import LayerNorm +from funasr.models.encoder.abs_encoder import AbsEncoder +import math +from funasr.models.transformer.utils.repeat import repeat +from funasr.models.transformer.utils.multi_layer_conv import FsmnFeedForward + + +class FsmnBlock(torch.nn.Module): + def __init__( + self, + n_feat, + dropout_rate, + kernel_size, + fsmn_shift=0, + ): + super().__init__() + self.dropout = nn.Dropout(p=dropout_rate) + self.fsmn_block = nn.Conv1d( + n_feat, n_feat, kernel_size, stride=1, padding=0, groups=n_feat, bias=False + ) + # padding + left_padding = (kernel_size - 1) // 2 + if fsmn_shift > 0: + left_padding = left_padding + fsmn_shift + right_padding = kernel_size - 1 - left_padding + self.pad_fn = nn.ConstantPad1d((left_padding, right_padding), 0.0) + + def forward(self, inputs, mask, mask_shfit_chunk=None): + b, t, d = inputs.size() + if mask is not None: + mask = torch.reshape(mask, (b, -1, 1)) + if mask_shfit_chunk is not None: + mask = mask * mask_shfit_chunk + + inputs = inputs * mask + x = inputs.transpose(1, 2) + x = self.pad_fn(x) + x = self.fsmn_block(x) + x = x.transpose(1, 2) + x = x + inputs + x = self.dropout(x) + return x * mask + + +class EncoderLayer(torch.nn.Module): + def __init__(self, in_size, size, feed_forward, fsmn_block, dropout_rate=0.0): + super().__init__() + self.in_size = in_size + self.size = size + self.ffn = feed_forward + self.memory = fsmn_block + self.dropout = nn.Dropout(dropout_rate) + + def forward( + self, xs_pad: torch.Tensor, mask: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + # xs_pad in Batch, Time, Dim + + context = self.ffn(xs_pad)[0] + memory = self.memory(context, mask) + + memory = self.dropout(memory) + if self.in_size == self.size: + return memory + xs_pad, mask + + return memory, mask + + +class FsmnEncoder(AbsEncoder): + """Encoder using Fsmn""" + + def __init__( + self, + in_units, + filter_size, + fsmn_num_layers, + dnn_num_layers, + num_memory_units=512, + ffn_inner_dim=2048, + dropout_rate=0.0, + shift=0, + position_encoder=None, + sample_rate=1, + out_units=None, + tf2torch_tensor_name_prefix_torch="post_net", + tf2torch_tensor_name_prefix_tf="EAND/post_net", + ): + """Initializes the parameters of the encoder. + + Args: + filter_size: the total order of memory block + fsmn_num_layers: The number of fsmn layers. + dnn_num_layers: The number of dnn layers + num_units: The number of memory units. + ffn_inner_dim: The number of units of the inner linear transformation + in the feed forward layer. + dropout_rate: The probability to drop units from the outputs. + shift: left padding, to control delay + position_encoder: The :class:`opennmt.layers.position.PositionEncoder` to + apply on inputs or ``None``. + """ + super(FsmnEncoder, self).__init__() + self.in_units = in_units + self.filter_size = filter_size + self.fsmn_num_layers = fsmn_num_layers + self.dnn_num_layers = dnn_num_layers + self.num_memory_units = num_memory_units + self.ffn_inner_dim = ffn_inner_dim + self.dropout_rate = dropout_rate + self.shift = shift + if not isinstance(shift, list): + self.shift = [shift for _ in range(self.fsmn_num_layers)] + self.sample_rate = sample_rate + if not isinstance(sample_rate, list): + self.sample_rate = [sample_rate for _ in range(self.fsmn_num_layers)] + self.position_encoder = position_encoder + self.dropout = nn.Dropout(dropout_rate) + self.out_units = out_units + self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch + self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf + + self.fsmn_layers = repeat( + self.fsmn_num_layers, + lambda lnum: EncoderLayer( + in_units if lnum == 0 else num_memory_units, + num_memory_units, + FsmnFeedForward( + in_units if lnum == 0 else num_memory_units, + ffn_inner_dim, + num_memory_units, + 1, + dropout_rate, + ), + FsmnBlock(num_memory_units, dropout_rate, filter_size, self.shift[lnum]), + ), + ) + + self.dnn_layers = repeat( + dnn_num_layers, + lambda lnum: FsmnFeedForward( + num_memory_units, + ffn_inner_dim, + num_memory_units, + 1, + dropout_rate, + ), + ) + if out_units is not None: + self.conv1d = nn.Conv1d(num_memory_units, out_units, 1, 1) + + def output_size(self) -> int: + return self.num_memory_units + + def forward( + self, xs_pad: torch.Tensor, ilens: torch.Tensor, prev_states: torch.Tensor = None + ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + inputs = xs_pad + if self.position_encoder is not None: + inputs = self.position_encoder(inputs) + + inputs = self.dropout(inputs) + masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device) + inputs = self.fsmn_layers(inputs, masks)[0] + inputs = self.dnn_layers(inputs)[0] + + if self.out_units is not None: + inputs = self.conv1d(inputs.transpose(1, 2)).transpose(1, 2) + + return inputs, ilens, None diff --git a/funasr/models/sond/encoder/resnet34_encoder.py b/funasr/models/sond/encoder/resnet34_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..ba577a3a445b47eb5fe61d4b1a4069864bafdbf0 --- /dev/null +++ b/funasr/models/sond/encoder/resnet34_encoder.py @@ -0,0 +1,462 @@ +import torch +from torch.nn import functional as F +from funasr.models.encoder.abs_encoder import AbsEncoder +from typing import Tuple, Optional +from funasr.models.pooling.statistic_pooling import statistic_pooling, windowed_statistic_pooling +from collections import OrderedDict +import logging +import numpy as np + + +class BasicLayer(torch.nn.Module): + + def __init__(self, in_filters: int, filters: int, stride: int, bn_momentum: float = 0.5): + + super().__init__() + self.stride = stride + self.in_filters = in_filters + self.filters = filters + + self.bn1 = torch.nn.BatchNorm2d(in_filters, eps=1e-3, momentum=bn_momentum, affine=True) + self.relu1 = torch.nn.ReLU() + self.conv1 = torch.nn.Conv2d(in_filters, filters, 3, stride, bias=False) + + self.bn2 = torch.nn.BatchNorm2d(filters, eps=1e-3, momentum=bn_momentum, affine=True) + self.relu2 = torch.nn.ReLU() + self.conv2 = torch.nn.Conv2d(filters, filters, 3, 1, bias=False) + + if in_filters != filters or stride > 1: + self.conv_sc = torch.nn.Conv2d(in_filters, filters, 1, stride, bias=False) + self.bn_sc = torch.nn.BatchNorm2d(filters, eps=1e-3, momentum=bn_momentum, affine=True) + + def proper_padding(self, x, stride): + # align padding mode to tf.layers.conv2d with padding_mod="same" + if stride == 1: + return F.pad(x, (1, 1, 1, 1), "constant", 0) + elif stride == 2: + h, w = x.size(2), x.size(3) + # (left, right, top, bottom) + return F.pad(x, (w % 2, 1, h % 2, 1), "constant", 0) + + def forward(self, xs_pad, ilens): + identity = xs_pad + if self.in_filters != self.filters or self.stride > 1: + identity = self.conv_sc(identity) + identity = self.bn_sc(identity) + + xs_pad = self.relu1(self.bn1(xs_pad)) + xs_pad = self.proper_padding(xs_pad, self.stride) + xs_pad = self.conv1(xs_pad) + + xs_pad = self.relu2(self.bn2(xs_pad)) + xs_pad = self.proper_padding(xs_pad, 1) + xs_pad = self.conv2(xs_pad) + + if self.stride == 2: + ilens = (ilens + 1) // self.stride + + return xs_pad + identity, ilens + + +class BasicBlock(torch.nn.Module): + def __init__(self, in_filters, filters, num_layer, stride, bn_momentum=0.5): + super().__init__() + self.num_layer = num_layer + + for i in range(num_layer): + layer = BasicLayer( + in_filters if i == 0 else filters, filters, stride if i == 0 else 1, bn_momentum + ) + self.add_module("layer_{}".format(i), layer) + + def forward(self, xs_pad, ilens): + + for i in range(self.num_layer): + xs_pad, ilens = self._modules["layer_{}".format(i)](xs_pad, ilens) + + return xs_pad, ilens + + +class ResNet34(AbsEncoder): + def __init__( + self, + input_size, + use_head_conv=True, + batchnorm_momentum=0.5, + use_head_maxpool=False, + num_nodes_pooling_layer=256, + layers_in_block=(3, 4, 6, 3), + filters_in_block=(32, 64, 128, 256), + ): + super(ResNet34, self).__init__() + + self.use_head_conv = use_head_conv + self.use_head_maxpool = use_head_maxpool + self.num_nodes_pooling_layer = num_nodes_pooling_layer + self.layers_in_block = layers_in_block + self.filters_in_block = filters_in_block + self.input_size = input_size + + pre_filters = filters_in_block[0] + if use_head_conv: + self.pre_conv = torch.nn.Conv2d( + 1, pre_filters, 3, 1, 1, bias=False, padding_mode="zeros" + ) + self.pre_conv_bn = torch.nn.BatchNorm2d( + pre_filters, eps=1e-3, momentum=batchnorm_momentum + ) + + if use_head_maxpool: + self.head_maxpool = torch.nn.MaxPool2d(3, 1, padding=1) + + for i in range(len(layers_in_block)): + if i == 0: + in_filters = pre_filters if self.use_head_conv else 1 + else: + in_filters = filters_in_block[i - 1] + + block = BasicBlock( + in_filters, + filters=filters_in_block[i], + num_layer=layers_in_block[i], + stride=1 if i == 0 else 2, + bn_momentum=batchnorm_momentum, + ) + self.add_module("block_{}".format(i), block) + + self.resnet0_dense = torch.nn.Conv2d(filters_in_block[-1], num_nodes_pooling_layer, 1) + self.resnet0_bn = torch.nn.BatchNorm2d( + num_nodes_pooling_layer, eps=1e-3, momentum=batchnorm_momentum + ) + + self.time_ds_ratio = 8 + + def output_size(self) -> int: + return self.num_nodes_pooling_layer + + def forward( + self, xs_pad: torch.Tensor, ilens: torch.Tensor, prev_states: torch.Tensor = None + ) -> Tuple[torch.Tensor, torch.Tensor]: + + features = xs_pad + assert ( + features.size(-1) == self.input_size + ), "Dimension of features {} doesn't match the input_size {}.".format( + features.size(-1), self.input_size + ) + features = torch.unsqueeze(features, dim=1) + if self.use_head_conv: + features = self.pre_conv(features) + features = self.pre_conv_bn(features) + features = F.relu(features) + + if self.use_head_maxpool: + features = self.head_maxpool(features) + + resnet_outs, resnet_out_lens = features, ilens + for i in range(len(self.layers_in_block)): + block = self._modules["block_{}".format(i)] + resnet_outs, resnet_out_lens = block(resnet_outs, resnet_out_lens) + + features = self.resnet0_dense(resnet_outs) + features = F.relu(features) + features = self.resnet0_bn(features) + + return features, resnet_out_lens + + +# Note: For training, this implement is not equivalent to tf because of the kernel_regularizer in tf.layers. +# TODO: implement kernel_regularizer in torch with munal loss addition or weigth_decay in the optimizer +class ResNet34_SP_L2Reg(AbsEncoder): + def __init__( + self, + input_size, + use_head_conv=True, + batchnorm_momentum=0.5, + use_head_maxpool=False, + num_nodes_pooling_layer=256, + layers_in_block=(3, 4, 6, 3), + filters_in_block=(32, 64, 128, 256), + tf2torch_tensor_name_prefix_torch="encoder", + tf2torch_tensor_name_prefix_tf="EAND/speech_encoder", + tf_train_steps=720000, + ): + super(ResNet34_SP_L2Reg, self).__init__() + + self.use_head_conv = use_head_conv + self.use_head_maxpool = use_head_maxpool + self.num_nodes_pooling_layer = num_nodes_pooling_layer + self.layers_in_block = layers_in_block + self.filters_in_block = filters_in_block + self.input_size = input_size + self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch + self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf + self.tf_train_steps = tf_train_steps + + pre_filters = filters_in_block[0] + if use_head_conv: + self.pre_conv = torch.nn.Conv2d( + 1, pre_filters, 3, 1, 1, bias=False, padding_mode="zeros" + ) + self.pre_conv_bn = torch.nn.BatchNorm2d( + pre_filters, eps=1e-3, momentum=batchnorm_momentum + ) + + if use_head_maxpool: + self.head_maxpool = torch.nn.MaxPool2d(3, 1, padding=1) + + for i in range(len(layers_in_block)): + if i == 0: + in_filters = pre_filters if self.use_head_conv else 1 + else: + in_filters = filters_in_block[i - 1] + + block = BasicBlock( + in_filters, + filters=filters_in_block[i], + num_layer=layers_in_block[i], + stride=1 if i == 0 else 2, + bn_momentum=batchnorm_momentum, + ) + self.add_module("block_{}".format(i), block) + + self.resnet0_dense = torch.nn.Conv1d( + filters_in_block[-1] * input_size // 8, num_nodes_pooling_layer, 1 + ) + self.resnet0_bn = torch.nn.BatchNorm1d( + num_nodes_pooling_layer, eps=1e-3, momentum=batchnorm_momentum + ) + + self.time_ds_ratio = 8 + + def output_size(self) -> int: + return self.num_nodes_pooling_layer + + def forward( + self, xs_pad: torch.Tensor, ilens: torch.Tensor, prev_states: torch.Tensor = None + ) -> Tuple[torch.Tensor, torch.Tensor]: + + features = xs_pad + assert ( + features.size(-1) == self.input_size + ), "Dimension of features {} doesn't match the input_size {}.".format( + features.size(-1), self.input_size + ) + features = torch.unsqueeze(features, dim=1) + if self.use_head_conv: + features = self.pre_conv(features) + features = self.pre_conv_bn(features) + features = F.relu(features) + + if self.use_head_maxpool: + features = self.head_maxpool(features) + + resnet_outs, resnet_out_lens = features, ilens + for i in range(len(self.layers_in_block)): + block = self._modules["block_{}".format(i)] + resnet_outs, resnet_out_lens = block(resnet_outs, resnet_out_lens) + + # B, C, T, F + bb, cc, tt, ff = resnet_outs.shape + resnet_outs = torch.reshape(resnet_outs.permute(0, 3, 1, 2), [bb, ff * cc, tt]) + features = self.resnet0_dense(resnet_outs) + features = F.relu(features) + features = self.resnet0_bn(features) + + return features, resnet_out_lens + + +class ResNet34Diar(ResNet34): + def __init__( + self, + input_size, + embedding_node="resnet1_dense", + use_head_conv=True, + batchnorm_momentum=0.5, + use_head_maxpool=False, + num_nodes_pooling_layer=256, + layers_in_block=(3, 4, 6, 3), + filters_in_block=(32, 64, 128, 256), + num_nodes_resnet1=256, + num_nodes_last_layer=256, + pooling_type="window_shift", + pool_size=20, + stride=1, + tf2torch_tensor_name_prefix_torch="encoder", + tf2torch_tensor_name_prefix_tf="seq2seq/speech_encoder", + ): + """ + Author: Speech Lab, Alibaba Group, China + SOND: Speaker Overlap-aware Neural Diarization for Multi-party Meeting Analysis + https://arxiv.org/abs/2211.10243 + """ + + super(ResNet34Diar, self).__init__( + input_size, + use_head_conv=use_head_conv, + batchnorm_momentum=batchnorm_momentum, + use_head_maxpool=use_head_maxpool, + num_nodes_pooling_layer=num_nodes_pooling_layer, + layers_in_block=layers_in_block, + filters_in_block=filters_in_block, + ) + + self.embedding_node = embedding_node + self.num_nodes_resnet1 = num_nodes_resnet1 + self.num_nodes_last_layer = num_nodes_last_layer + self.pooling_type = pooling_type + self.pool_size = pool_size + self.stride = stride + self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch + self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf + + self.resnet1_dense = torch.nn.Linear(num_nodes_pooling_layer * 2, num_nodes_resnet1) + self.resnet1_bn = torch.nn.BatchNorm1d( + num_nodes_resnet1, eps=1e-3, momentum=batchnorm_momentum + ) + + self.resnet2_dense = torch.nn.Linear(num_nodes_resnet1, num_nodes_last_layer) + self.resnet2_bn = torch.nn.BatchNorm1d( + num_nodes_last_layer, eps=1e-3, momentum=batchnorm_momentum + ) + + def output_size(self) -> int: + if self.embedding_node.startswith("resnet1"): + return self.num_nodes_resnet1 + elif self.embedding_node.startswith("resnet2"): + return self.num_nodes_last_layer + + return self.num_nodes_pooling_layer + + def forward( + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + prev_states: torch.Tensor = None, + ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + + endpoints = OrderedDict() + res_out, ilens = super().forward(xs_pad, ilens) + endpoints["resnet0_bn"] = res_out + if self.pooling_type == "frame_gsp": + features = statistic_pooling(res_out, ilens, (3,)) + else: + features, ilens = windowed_statistic_pooling( + res_out, ilens, (2, 3), self.pool_size, self.stride + ) + features = features.transpose(1, 2) + endpoints["pooling"] = features + + features = self.resnet1_dense(features) + endpoints["resnet1_dense"] = features + features = F.relu(features) + endpoints["resnet1_relu"] = features + features = self.resnet1_bn(features.transpose(1, 2)).transpose(1, 2) + endpoints["resnet1_bn"] = features + + features = self.resnet2_dense(features) + endpoints["resnet2_dense"] = features + features = F.relu(features) + endpoints["resnet2_relu"] = features + features = self.resnet2_bn(features.transpose(1, 2)).transpose(1, 2) + endpoints["resnet2_bn"] = features + + return endpoints[self.embedding_node], ilens, None + + +class ResNet34SpL2RegDiar(ResNet34_SP_L2Reg): + def __init__( + self, + input_size, + embedding_node="resnet1_dense", + use_head_conv=True, + batchnorm_momentum=0.5, + use_head_maxpool=False, + num_nodes_pooling_layer=256, + layers_in_block=(3, 4, 6, 3), + filters_in_block=(32, 64, 128, 256), + num_nodes_resnet1=256, + num_nodes_last_layer=256, + pooling_type="window_shift", + pool_size=20, + stride=1, + tf2torch_tensor_name_prefix_torch="encoder", + tf2torch_tensor_name_prefix_tf="seq2seq/speech_encoder", + ): + """ + Author: Speech Lab, Alibaba Group, China + TOLD: A Novel Two-Stage Overlap-Aware Framework for Speaker Diarization + https://arxiv.org/abs/2303.05397 + """ + + super(ResNet34SpL2RegDiar, self).__init__( + input_size, + use_head_conv=use_head_conv, + batchnorm_momentum=batchnorm_momentum, + use_head_maxpool=use_head_maxpool, + num_nodes_pooling_layer=num_nodes_pooling_layer, + layers_in_block=layers_in_block, + filters_in_block=filters_in_block, + ) + + self.embedding_node = embedding_node + self.num_nodes_resnet1 = num_nodes_resnet1 + self.num_nodes_last_layer = num_nodes_last_layer + self.pooling_type = pooling_type + self.pool_size = pool_size + self.stride = stride + self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch + self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf + + self.resnet1_dense = torch.nn.Linear(num_nodes_pooling_layer * 2, num_nodes_resnet1) + self.resnet1_bn = torch.nn.BatchNorm1d( + num_nodes_resnet1, eps=1e-3, momentum=batchnorm_momentum + ) + + self.resnet2_dense = torch.nn.Linear(num_nodes_resnet1, num_nodes_last_layer) + self.resnet2_bn = torch.nn.BatchNorm1d( + num_nodes_last_layer, eps=1e-3, momentum=batchnorm_momentum + ) + + def output_size(self) -> int: + if self.embedding_node.startswith("resnet1"): + return self.num_nodes_resnet1 + elif self.embedding_node.startswith("resnet2"): + return self.num_nodes_last_layer + + return self.num_nodes_pooling_layer + + def forward( + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + prev_states: torch.Tensor = None, + ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + + endpoints = OrderedDict() + res_out, ilens = super().forward(xs_pad, ilens) + endpoints["resnet0_bn"] = res_out + if self.pooling_type == "frame_gsp": + features = statistic_pooling(res_out, ilens, (2,)) + else: + features, ilens = windowed_statistic_pooling( + res_out, ilens, (2,), self.pool_size, self.stride + ) + features = features.transpose(1, 2) + endpoints["pooling"] = features + + features = self.resnet1_dense(features) + endpoints["resnet1_dense"] = features + features = F.relu(features) + endpoints["resnet1_relu"] = features + features = self.resnet1_bn(features.transpose(1, 2)).transpose(1, 2) + endpoints["resnet1_bn"] = features + + features = self.resnet2_dense(features) + endpoints["resnet2_dense"] = features + features = F.relu(features) + endpoints["resnet2_relu"] = features + features = self.resnet2_bn(features.transpose(1, 2)).transpose(1, 2) + endpoints["resnet2_bn"] = features + + return endpoints[self.embedding_node], ilens, None diff --git a/funasr/models/sond/encoder/self_attention_encoder.py b/funasr/models/sond/encoder/self_attention_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..2e979b1be95509a083e159e6c80c84e414fa08cb --- /dev/null +++ b/funasr/models/sond/encoder/self_attention_encoder.py @@ -0,0 +1,333 @@ +from typing import List +from typing import Optional +from typing import Sequence +from typing import Tuple +from typing import Union +import logging +import torch +import torch.nn as nn +from funasr.models.scama.chunk_utilis import overlap_chunk +import numpy as np +from funasr.models.transformer.utils.nets_utils import make_pad_mask +from funasr.models.sond.attention import MultiHeadSelfAttention +from funasr.models.transformer.embedding import SinusoidalPositionEncoder +from funasr.models.transformer.layer_norm import LayerNorm +from funasr.models.transformer.utils.multi_layer_conv import Conv1dLinear +from funasr.models.transformer.utils.multi_layer_conv import MultiLayeredConv1d +from funasr.models.transformer.positionwise_feed_forward import ( + PositionwiseFeedForward, # noqa: H301 +) +from funasr.models.transformer.utils.repeat import repeat +from funasr.models.transformer.utils.subsampling import Conv2dSubsampling +from funasr.models.transformer.utils.subsampling import Conv2dSubsampling2 +from funasr.models.transformer.utils.subsampling import Conv2dSubsampling6 +from funasr.models.transformer.utils.subsampling import Conv2dSubsampling8 +from funasr.models.transformer.utils.subsampling import TooShortUttError +from funasr.models.transformer.utils.subsampling import check_short_utt +from funasr.models.ctc import CTC +from funasr.models.encoder.abs_encoder import AbsEncoder + + +class EncoderLayer(nn.Module): + def __init__( + self, + in_size, + size, + self_attn, + feed_forward, + dropout_rate, + normalize_before=True, + concat_after=False, + stochastic_depth_rate=0.0, + ): + """Construct an EncoderLayer object.""" + super(EncoderLayer, self).__init__() + self.self_attn = self_attn + self.feed_forward = feed_forward + self.norm1 = LayerNorm(in_size) + self.norm2 = LayerNorm(size) + self.dropout = nn.Dropout(dropout_rate) + self.in_size = in_size + self.size = size + self.normalize_before = normalize_before + self.concat_after = concat_after + if self.concat_after: + self.concat_linear = nn.Linear(size + size, size) + self.stochastic_depth_rate = stochastic_depth_rate + self.dropout_rate = dropout_rate + + def forward(self, x, mask, cache=None, mask_att_chunk_encoder=None): + """Compute encoded features. + + Args: + x_input (torch.Tensor): Input tensor (#batch, time, size). + mask (torch.Tensor): Mask tensor for the input (#batch, time). + cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size). + + Returns: + torch.Tensor: Output tensor (#batch, time, size). + torch.Tensor: Mask tensor (#batch, time). + + """ + skip_layer = False + # with stochastic depth, residual connection `x + f(x)` becomes + # `x <- x + 1 / (1 - p) * f(x)` at training time. + stoch_layer_coeff = 1.0 + if self.training and self.stochastic_depth_rate > 0: + skip_layer = torch.rand(1).item() < self.stochastic_depth_rate + stoch_layer_coeff = 1.0 / (1 - self.stochastic_depth_rate) + + if skip_layer: + if cache is not None: + x = torch.cat([cache, x], dim=1) + return x, mask + + residual = x + if self.normalize_before: + x = self.norm1(x) + + if self.concat_after: + x_concat = torch.cat( + (x, self.self_attn(x, mask, mask_att_chunk_encoder=mask_att_chunk_encoder)), dim=-1 + ) + if self.in_size == self.size: + x = residual + stoch_layer_coeff * self.concat_linear(x_concat) + else: + x = stoch_layer_coeff * self.concat_linear(x_concat) + else: + if self.in_size == self.size: + x = residual + stoch_layer_coeff * self.dropout( + self.self_attn(x, mask, mask_att_chunk_encoder=mask_att_chunk_encoder) + ) + else: + x = stoch_layer_coeff * self.dropout( + self.self_attn(x, mask, mask_att_chunk_encoder=mask_att_chunk_encoder) + ) + if not self.normalize_before: + x = self.norm1(x) + + residual = x + if self.normalize_before: + x = self.norm2(x) + x = residual + stoch_layer_coeff * self.dropout(self.feed_forward(x)) + if not self.normalize_before: + x = self.norm2(x) + + return x, mask, cache, mask_att_chunk_encoder + + +class SelfAttentionEncoder(AbsEncoder): + """ + Author: Speech Lab of DAMO Academy, Alibaba Group + Self attention encoder in OpenNMT framework + """ + + def __init__( + self, + input_size: int, + output_size: int = 256, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + attention_dropout_rate: float = 0.0, + input_layer: Optional[str] = "conv2d", + pos_enc_class=SinusoidalPositionEncoder, + normalize_before: bool = True, + concat_after: bool = False, + positionwise_layer_type: str = "linear", + positionwise_conv_kernel_size: int = 1, + padding_idx: int = -1, + interctc_layer_idx: List[int] = [], + interctc_use_conditioning: bool = False, + tf2torch_tensor_name_prefix_torch: str = "encoder", + tf2torch_tensor_name_prefix_tf: str = "seq2seq/encoder", + out_units=None, + ): + super().__init__() + self._output_size = output_size + + if input_layer == "linear": + self.embed = torch.nn.Sequential( + torch.nn.Linear(input_size, output_size), + torch.nn.LayerNorm(output_size), + torch.nn.Dropout(dropout_rate), + torch.nn.ReLU(), + pos_enc_class(output_size, positional_dropout_rate), + ) + elif input_layer == "conv2d": + self.embed = Conv2dSubsampling(input_size, output_size, dropout_rate) + elif input_layer == "conv2d2": + self.embed = Conv2dSubsampling2(input_size, output_size, dropout_rate) + elif input_layer == "conv2d6": + self.embed = Conv2dSubsampling6(input_size, output_size, dropout_rate) + elif input_layer == "conv2d8": + self.embed = Conv2dSubsampling8(input_size, output_size, dropout_rate) + elif input_layer == "embed": + self.embed = torch.nn.Sequential( + torch.nn.Embedding(input_size, output_size, padding_idx=padding_idx), + SinusoidalPositionEncoder(), + ) + elif input_layer is None: + if input_size == output_size: + self.embed = None + else: + self.embed = torch.nn.Linear(input_size, output_size) + elif input_layer == "pe": + self.embed = SinusoidalPositionEncoder() + elif input_layer == "null": + self.embed = None + else: + raise ValueError("unknown input_layer: " + input_layer) + self.normalize_before = normalize_before + if positionwise_layer_type == "linear": + positionwise_layer = PositionwiseFeedForward + positionwise_layer_args = ( + output_size, + linear_units, + dropout_rate, + ) + elif positionwise_layer_type == "conv1d": + positionwise_layer = MultiLayeredConv1d + positionwise_layer_args = ( + output_size, + linear_units, + positionwise_conv_kernel_size, + dropout_rate, + ) + elif positionwise_layer_type == "conv1d-linear": + positionwise_layer = Conv1dLinear + positionwise_layer_args = ( + output_size, + linear_units, + positionwise_conv_kernel_size, + dropout_rate, + ) + else: + raise NotImplementedError("Support only linear or conv1d.") + + self.encoders = repeat( + num_blocks, + lambda lnum: ( + EncoderLayer( + output_size, + output_size, + MultiHeadSelfAttention( + attention_heads, + output_size, + output_size, + attention_dropout_rate, + ), + positionwise_layer(*positionwise_layer_args), + dropout_rate, + normalize_before, + concat_after, + ) + if lnum > 0 + else EncoderLayer( + input_size, + output_size, + MultiHeadSelfAttention( + attention_heads, + input_size if input_layer == "pe" or input_layer == "null" else output_size, + output_size, + attention_dropout_rate, + ), + positionwise_layer(*positionwise_layer_args), + dropout_rate, + normalize_before, + concat_after, + ) + ), + ) + if self.normalize_before: + self.after_norm = LayerNorm(output_size) + + self.interctc_layer_idx = interctc_layer_idx + if len(interctc_layer_idx) > 0: + assert 0 < min(interctc_layer_idx) and max(interctc_layer_idx) < num_blocks + self.interctc_use_conditioning = interctc_use_conditioning + self.conditioning_layer = None + self.dropout = nn.Dropout(dropout_rate) + self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch + self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf + self.out_units = out_units + if out_units is not None: + self.output_linear = nn.Linear(output_size, out_units) + + def output_size(self) -> int: + return self._output_size + + def forward( + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + prev_states: torch.Tensor = None, + ctc: CTC = None, + ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Embed positions in tensor. + + Args: + xs_pad: input tensor (B, L, D) + ilens: input length (B) + prev_states: Not to be used now. + Returns: + position embedded tensor and mask + """ + masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device) + xs_pad = xs_pad * self.output_size() ** 0.5 + if self.embed is None: + xs_pad = xs_pad + elif ( + isinstance(self.embed, Conv2dSubsampling) + or isinstance(self.embed, Conv2dSubsampling2) + or isinstance(self.embed, Conv2dSubsampling6) + or isinstance(self.embed, Conv2dSubsampling8) + ): + short_status, limit_size = check_short_utt(self.embed, xs_pad.size(1)) + if short_status: + raise TooShortUttError( + f"has {xs_pad.size(1)} frames and is too short for subsampling " + + f"(it needs more than {limit_size} frames), return empty results", + xs_pad.size(1), + limit_size, + ) + xs_pad, masks = self.embed(xs_pad, masks) + else: + xs_pad = self.embed(xs_pad) + + xs_pad = self.dropout(xs_pad) + # encoder_outs = self.encoders0(xs_pad, masks) + # xs_pad, masks = encoder_outs[0], encoder_outs[1] + intermediate_outs = [] + if len(self.interctc_layer_idx) == 0: + encoder_outs = self.encoders(xs_pad, masks) + xs_pad, masks = encoder_outs[0], encoder_outs[1] + else: + for layer_idx, encoder_layer in enumerate(self.encoders): + encoder_outs = encoder_layer(xs_pad, masks) + xs_pad, masks = encoder_outs[0], encoder_outs[1] + + if layer_idx + 1 in self.interctc_layer_idx: + encoder_out = xs_pad + + # intermediate outputs are also normalized + if self.normalize_before: + encoder_out = self.after_norm(encoder_out) + + intermediate_outs.append((layer_idx + 1, encoder_out)) + + if self.interctc_use_conditioning: + ctc_out = ctc.softmax(encoder_out) + xs_pad = xs_pad + self.conditioning_layer(ctc_out) + + if self.normalize_before: + xs_pad = self.after_norm(xs_pad) + + if self.out_units is not None: + xs_pad = self.output_linear(xs_pad) + olens = masks.squeeze(1).sum(1) + if len(intermediate_outs) > 0: + return (xs_pad, intermediate_outs), olens, None + return xs_pad, olens, None diff --git a/funasr/models/sond/label_aggregation.py b/funasr/models/sond/label_aggregation.py new file mode 100644 index 0000000000000000000000000000000000000000..dddc4c4bf4dbadd7a2bf5ddb3491ba69bb04b615 --- /dev/null +++ b/funasr/models/sond/label_aggregation.py @@ -0,0 +1,113 @@ +import torch +from typing import Optional +from typing import Tuple +from torch.nn import functional as F +from funasr.models.transformer.utils.nets_utils import make_pad_mask + + +class LabelAggregate(torch.nn.Module): + def __init__( + self, + win_length: int = 512, + hop_length: int = 128, + center: bool = True, + ): + super().__init__() + + self.win_length = win_length + self.hop_length = hop_length + self.center = center + + def extra_repr(self): + return ( + f"win_length={self.win_length}, " + f"hop_length={self.hop_length}, " + f"center={self.center}, " + ) + + def forward( + self, input: torch.Tensor, ilens: torch.Tensor = None + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """LabelAggregate forward function. + + Args: + input: (Batch, Nsamples, Label_dim) + ilens: (Batch) + Returns: + output: (Batch, Frames, Label_dim) + + """ + bs = input.size(0) + max_length = input.size(1) + label_dim = input.size(2) + + # NOTE(jiatong): + # The default behaviour of label aggregation is compatible with + # torch.stft about framing and padding. + + # Step1: center padding + if self.center: + pad = self.win_length // 2 + max_length = max_length + 2 * pad + input = torch.nn.functional.pad(input, (0, 0, pad, pad), "constant", 0) + input[:, :pad, :] = input[:, pad : (2 * pad), :] + input[:, (max_length - pad) : max_length, :] = input[ + :, (max_length - 2 * pad) : (max_length - pad), : + ] + nframe = (max_length - self.win_length) // self.hop_length + 1 + + # Step2: framing + output = input.as_strided( + (bs, nframe, self.win_length, label_dim), + (max_length * label_dim, self.hop_length * label_dim, label_dim, 1), + ) + + # Step3: aggregate label + output = torch.gt(output.sum(dim=2, keepdim=False), self.win_length // 2) + output = output.float() + + # Step4: process lengths + if ilens is not None: + if self.center: + pad = self.win_length // 2 + ilens = ilens + 2 * pad + + olens = (ilens - self.win_length) // self.hop_length + 1 + output.masked_fill_(make_pad_mask(olens, output, 1), 0.0) + else: + olens = None + + return output.to(input.dtype), olens + + +class LabelAggregateMaxPooling(torch.nn.Module): + def __init__( + self, + hop_length: int = 8, + ): + super().__init__() + + self.hop_length = hop_length + + def extra_repr(self): + return f"hop_length={self.hop_length}, " + + def forward( + self, input: torch.Tensor, ilens: torch.Tensor = None + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """LabelAggregate forward function. + + Args: + input: (Batch, Nsamples, Label_dim) + ilens: (Batch) + Returns: + output: (Batch, Frames, Label_dim) + + """ + + output = F.max_pool1d(input.transpose(1, 2), self.hop_length, self.hop_length).transpose( + 1, 2 + ) + olens = ilens // self.hop_length + + return output.to(input.dtype), olens diff --git a/funasr/models/sond/pooling/__init__.py b/funasr/models/sond/pooling/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/models/sond/pooling/pooling_layers.py b/funasr/models/sond/pooling/pooling_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..0a426a94fc701df0c9aba816792df889dfef53c4 --- /dev/null +++ b/funasr/models/sond/pooling/pooling_layers.py @@ -0,0 +1,107 @@ +# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +""" This implementation is adapted from https://github.com/wenet-e2e/wespeaker.""" + +import torch +import torch.nn as nn + + +class TAP(nn.Module): + """ + Temporal average pooling, only first-order mean is considered + """ + + def __init__(self, **kwargs): + super(TAP, self).__init__() + + def forward(self, x): + pooling_mean = x.mean(dim=-1) + # To be compatable with 2D input + pooling_mean = pooling_mean.flatten(start_dim=1) + return pooling_mean + + +class TSDP(nn.Module): + """ + Temporal standard deviation pooling, only second-order std is considered + """ + + def __init__(self, **kwargs): + super(TSDP, self).__init__() + + def forward(self, x): + # The last dimension is the temporal axis + pooling_std = torch.sqrt(torch.var(x, dim=-1) + 1e-8) + pooling_std = pooling_std.flatten(start_dim=1) + return pooling_std + + +class TSTP(nn.Module): + """ + Temporal statistics pooling, concatenate mean and std, which is used in + x-vector + Comment: simple concatenation can not make full use of both statistics + """ + + def __init__(self, **kwargs): + super(TSTP, self).__init__() + + def forward(self, x): + # The last dimension is the temporal axis + pooling_mean = x.mean(dim=-1) + pooling_std = torch.sqrt(torch.var(x, dim=-1) + 1e-8) + pooling_mean = pooling_mean.flatten(start_dim=1) + pooling_std = pooling_std.flatten(start_dim=1) + + stats = torch.cat((pooling_mean, pooling_std), 1) + return stats + + +class ASTP(nn.Module): + """Attentive statistics pooling: Channel- and context-dependent + statistics pooling, first used in ECAPA_TDNN. + """ + + def __init__(self, in_dim, bottleneck_dim=128, global_context_att=False): + super(ASTP, self).__init__() + self.global_context_att = global_context_att + + # Use Conv1d with stride == 1 rather than Linear, then we don't + # need to transpose inputs. + if global_context_att: + self.linear1 = nn.Conv1d( + in_dim * 3, bottleneck_dim, kernel_size=1 + ) # equals W and b in the paper + else: + self.linear1 = nn.Conv1d( + in_dim, bottleneck_dim, kernel_size=1 + ) # equals W and b in the paper + self.linear2 = nn.Conv1d( + bottleneck_dim, in_dim, kernel_size=1 + ) # equals V and k in the paper + + def forward(self, x): + """ + x: a 3-dimensional tensor in tdnn-based architecture (B,F,T) + or a 4-dimensional tensor in resnet architecture (B,C,F,T) + 0-dim: batch-dimension, last-dim: time-dimension (frame-dimension) + """ + if len(x.shape) == 4: + x = x.reshape(x.shape[0], x.shape[1] * x.shape[2], x.shape[3]) + assert len(x.shape) == 3 + + if self.global_context_att: + context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x) + context_std = torch.sqrt(torch.var(x, dim=-1, keepdim=True) + 1e-10).expand_as(x) + x_in = torch.cat((x, context_mean, context_std), dim=1) + else: + x_in = x + + # DON'T use ReLU here! ReLU may be hard to converge. + alpha = torch.tanh(self.linear1(x_in)) # alpha = F.relu(self.linear1(x_in)) + alpha = torch.softmax(self.linear2(alpha), dim=2) + mean = torch.sum(alpha * x, dim=2) + var = torch.sum(alpha * (x**2), dim=2) - mean**2 + std = torch.sqrt(var.clamp(min=1e-10)) + return torch.cat([mean, std], dim=1) diff --git a/funasr/models/sond/pooling/statistic_pooling.py b/funasr/models/sond/pooling/statistic_pooling.py new file mode 100644 index 0000000000000000000000000000000000000000..91db2393e9d9a707ce9e0d8ebbaf4ea7d2f813ae --- /dev/null +++ b/funasr/models/sond/pooling/statistic_pooling.py @@ -0,0 +1,98 @@ +import torch +from typing import Tuple +from typing import Union +from funasr.models.transformer.utils.nets_utils import make_non_pad_mask +from torch.nn import functional as F +import math + +VAR2STD_EPSILON = 1e-12 + + +class StatisticPooling(torch.nn.Module): + def __init__(self, pooling_dim: Union[int, Tuple] = 2, eps=1e-12): + super(StatisticPooling, self).__init__() + if isinstance(pooling_dim, int): + pooling_dim = (pooling_dim,) + self.pooling_dim = pooling_dim + self.eps = eps + + def forward(self, xs_pad, ilens=None): + # xs_pad in (Batch, Channel, Time, Frequency) + + if ilens is None: + masks = torch.ones_like(xs_pad).to(xs_pad) + else: + masks = make_non_pad_mask(ilens, xs_pad, length_dim=2).to(xs_pad) + mean = torch.sum(xs_pad, dim=self.pooling_dim, keepdim=True) / torch.sum( + masks, dim=self.pooling_dim, keepdim=True + ) + squared_difference = torch.pow(xs_pad - mean, 2.0) + variance = torch.sum(squared_difference, dim=self.pooling_dim, keepdim=True) / torch.sum( + masks, dim=self.pooling_dim, keepdim=True + ) + for i in reversed(self.pooling_dim): + mean, variance = torch.squeeze(mean, dim=i), torch.squeeze(variance, dim=i) + + mask = torch.less_equal(variance, self.eps).float() + variance = (1.0 - mask) * variance + mask * self.eps + stddev = torch.sqrt(variance) + + stat_pooling = torch.cat([mean, stddev], dim=1) + + return stat_pooling + + +def statistic_pooling( + xs_pad: torch.Tensor, ilens: torch.Tensor = None, pooling_dim: Tuple = (2, 3) +) -> torch.Tensor: + # xs_pad in (Batch, Channel, Time, Frequency) + + if ilens is None: + seq_mask = torch.ones_like(xs_pad).to(xs_pad) + else: + seq_mask = make_non_pad_mask(ilens, xs_pad, length_dim=2).to(xs_pad) + mean = torch.sum(xs_pad, dim=pooling_dim, keepdim=True) / torch.sum( + seq_mask, dim=pooling_dim, keepdim=True + ) + squared_difference = torch.pow(xs_pad - mean, 2.0) + variance = torch.sum(squared_difference, dim=pooling_dim, keepdim=True) / torch.sum( + seq_mask, dim=pooling_dim, keepdim=True + ) + for i in reversed(pooling_dim): + mean, variance = torch.squeeze(mean, dim=i), torch.squeeze(variance, dim=i) + + value_mask = torch.less_equal(variance, VAR2STD_EPSILON).float() + variance = (1.0 - value_mask) * variance + value_mask * VAR2STD_EPSILON + stddev = torch.sqrt(variance) + + stat_pooling = torch.cat([mean, stddev], dim=1) + + return stat_pooling + + +def windowed_statistic_pooling( + xs_pad: torch.Tensor, + ilens: torch.Tensor = None, + pooling_dim: Tuple = (2, 3), + pooling_size: int = 20, + pooling_stride: int = 1, +) -> Tuple[torch.Tensor, int]: + # xs_pad in (Batch, Channel, Time, Frequency) + + tt = xs_pad.shape[2] + num_chunk = int(math.ceil(tt / pooling_stride)) + pad = pooling_size // 2 + if len(xs_pad.shape) == 4: + features = F.pad(xs_pad, (0, 0, pad, pad), "replicate") + else: + features = F.pad(xs_pad, (pad, pad), "replicate") + stat_list = [] + + for i in range(num_chunk): + # B x C + st, ed = i * pooling_stride, i * pooling_stride + pooling_size + stat = statistic_pooling(features[:, :, st:ed], pooling_dim=pooling_dim) + stat_list.append(stat.unsqueeze(2)) + + # B x C x T + return torch.cat(stat_list, dim=2), ilens / pooling_stride diff --git a/funasr/models/sond/sv_decoder.py b/funasr/models/sond/sv_decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..7d8484db980b617be209d41038c6a2d688a588e3 --- /dev/null +++ b/funasr/models/sond/sv_decoder.py @@ -0,0 +1,41 @@ +import torch +from torch.nn import functional as F +from funasr.models.decoder.abs_decoder import AbsDecoder + + +class DenseDecoder(AbsDecoder): + def __init__( + self, + vocab_size, + encoder_output_size, + num_nodes_resnet1: int = 256, + num_nodes_last_layer: int = 256, + batchnorm_momentum: float = 0.5, + ): + super(DenseDecoder, self).__init__() + self.resnet1_dense = torch.nn.Linear(encoder_output_size, num_nodes_resnet1) + self.resnet1_bn = torch.nn.BatchNorm1d( + num_nodes_resnet1, eps=1e-3, momentum=batchnorm_momentum + ) + + self.resnet2_dense = torch.nn.Linear(num_nodes_resnet1, num_nodes_last_layer) + self.resnet2_bn = torch.nn.BatchNorm1d( + num_nodes_last_layer, eps=1e-3, momentum=batchnorm_momentum + ) + + self.output_dense = torch.nn.Linear(num_nodes_last_layer, vocab_size, bias=False) + + def forward(self, features): + embeddings = {} + features = self.resnet1_dense(features) + embeddings["resnet1_dense"] = features + features = F.relu(features) + features = self.resnet1_bn(features) + + features = self.resnet2_dense(features) + embeddings["resnet2_dense"] = features + features = F.relu(features) + features = self.resnet2_bn(features) + + features = self.output_dense(features) + return features, embeddings diff --git a/funasr/models/specaug/__init__.py b/funasr/models/specaug/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/models/specaug/mask_along_axis.py b/funasr/models/specaug/mask_along_axis.py new file mode 100644 index 0000000000000000000000000000000000000000..d4f237b3d306237e55393a6bf65f984b3b789932 --- /dev/null +++ b/funasr/models/specaug/mask_along_axis.py @@ -0,0 +1,336 @@ +import math +import torch +from typing import Sequence +from typing import Union + + +def mask_along_axis( + spec: torch.Tensor, + spec_lengths: torch.Tensor, + mask_width_range: Sequence[int] = (0, 30), + dim: int = 1, + num_mask: int = 2, + replace_with_zero: bool = True, +): + """Apply mask along the specified direction. + + Args: + spec: (Batch, Length, Freq) + spec_lengths: (Length): Not using lengths in this implementation + mask_width_range: Select the width randomly between this range + """ + + org_size = spec.size() + if spec.dim() == 4: + # spec: (Batch, Channel, Length, Freq) -> (Batch * Channel, Length, Freq) + spec = spec.view(-1, spec.size(2), spec.size(3)) + + B = spec.shape[0] + # D = Length or Freq + D = spec.shape[dim] + # mask_length: (B, num_mask, 1) + mask_length = torch.randint( + mask_width_range[0], + mask_width_range[1], + (B, num_mask), + device=spec.device, + ).unsqueeze(2) + + # mask_pos: (B, num_mask, 1) + mask_pos = torch.randint( + 0, max(1, D - mask_length.max()), (B, num_mask), device=spec.device + ).unsqueeze(2) + + # aran: (1, 1, D) + aran = torch.arange(D, device=spec.device)[None, None, :] + # mask: (Batch, num_mask, D) + mask = (mask_pos <= aran) * (aran < (mask_pos + mask_length)) + # Multiply masks: (Batch, num_mask, D) -> (Batch, D) + mask = mask.any(dim=1) + if dim == 1: + # mask: (Batch, Length, 1) + mask = mask.unsqueeze(2) + elif dim == 2: + # mask: (Batch, 1, Freq) + mask = mask.unsqueeze(1) + + if replace_with_zero: + value = 0.0 + else: + value = spec.mean() + + if spec.requires_grad: + spec = spec.masked_fill(mask, value) + else: + spec = spec.masked_fill_(mask, value) + spec = spec.view(*org_size) + return spec, spec_lengths + + +def mask_along_axis_lfr( + spec: torch.Tensor, + spec_lengths: torch.Tensor, + mask_width_range: Sequence[int] = (0, 30), + dim: int = 1, + num_mask: int = 2, + replace_with_zero: bool = True, + lfr_rate: int = 1, +): + """Apply mask along the specified direction. + + Args: + spec: (Batch, Length, Freq) + spec_lengths: (Length): Not using lengths in this implementation + mask_width_range: Select the width randomly between this range + lfr_rate：low frame rate + """ + + org_size = spec.size() + if spec.dim() == 4: + # spec: (Batch, Channel, Length, Freq) -> (Batch * Channel, Length, Freq) + spec = spec.view(-1, spec.size(2), spec.size(3)) + + B = spec.shape[0] + # D = Length or Freq + D = spec.shape[dim] // lfr_rate + # mask_length: (B, num_mask, 1) + mask_length = torch.randint( + mask_width_range[0], + mask_width_range[1], + (B, num_mask), + device=spec.device, + ).unsqueeze(2) + if lfr_rate > 1: + mask_length = mask_length.repeat(1, lfr_rate, 1) + # mask_pos: (B, num_mask, 1) + mask_pos = torch.randint( + 0, max(1, D - mask_length.max()), (B, num_mask), device=spec.device + ).unsqueeze(2) + if lfr_rate > 1: + mask_pos_raw = mask_pos.clone() + mask_pos = torch.zeros((B, 0, 1), device=spec.device, dtype=torch.int32) + for i in range(lfr_rate): + mask_pos_i = mask_pos_raw + D * i + mask_pos = torch.cat((mask_pos, mask_pos_i), dim=1) + # aran: (1, 1, D) + D = spec.shape[dim] + aran = torch.arange(D, device=spec.device)[None, None, :] + # mask: (Batch, num_mask, D) + mask = (mask_pos <= aran) * (aran < (mask_pos + mask_length)) + # Multiply masks: (Batch, num_mask, D) -> (Batch, D) + mask = mask.any(dim=1) + if dim == 1: + # mask: (Batch, Length, 1) + mask = mask.unsqueeze(2) + elif dim == 2: + # mask: (Batch, 1, Freq) + mask = mask.unsqueeze(1) + + if replace_with_zero: + value = 0.0 + else: + value = spec.mean() + + if spec.requires_grad: + spec = spec.masked_fill(mask, value) + else: + spec = spec.masked_fill_(mask, value) + spec = spec.view(*org_size) + return spec, spec_lengths + + +class MaskAlongAxis(torch.nn.Module): + def __init__( + self, + mask_width_range: Union[int, Sequence[int]] = (0, 30), + num_mask: int = 2, + dim: Union[int, str] = "time", + replace_with_zero: bool = True, + ): + if isinstance(mask_width_range, int): + mask_width_range = (0, mask_width_range) + if len(mask_width_range) != 2: + raise TypeError( + f"mask_width_range must be a tuple of int and int values: " f"{mask_width_range}", + ) + + assert mask_width_range[1] > mask_width_range[0] + if isinstance(dim, str): + if dim == "time": + dim = 1 + elif dim == "freq": + dim = 2 + else: + raise ValueError("dim must be int, 'time' or 'freq'") + if dim == 1: + self.mask_axis = "time" + elif dim == 2: + self.mask_axis = "freq" + else: + self.mask_axis = "unknown" + + super().__init__() + self.mask_width_range = mask_width_range + self.num_mask = num_mask + self.dim = dim + self.replace_with_zero = replace_with_zero + + def extra_repr(self): + return ( + f"mask_width_range={self.mask_width_range}, " + f"num_mask={self.num_mask}, axis={self.mask_axis}" + ) + + def forward(self, spec: torch.Tensor, spec_lengths: torch.Tensor = None): + """Forward function. + + Args: + spec: (Batch, Length, Freq) + """ + + return mask_along_axis( + spec, + spec_lengths, + mask_width_range=self.mask_width_range, + dim=self.dim, + num_mask=self.num_mask, + replace_with_zero=self.replace_with_zero, + ) + + +class MaskAlongAxisVariableMaxWidth(torch.nn.Module): + """Mask input spec along a specified axis with variable maximum width. + + Formula: + max_width = max_width_ratio * seq_len + """ + + def __init__( + self, + mask_width_ratio_range: Union[float, Sequence[float]] = (0.0, 0.05), + num_mask: int = 2, + dim: Union[int, str] = "time", + replace_with_zero: bool = True, + ): + if isinstance(mask_width_ratio_range, float): + mask_width_ratio_range = (0.0, mask_width_ratio_range) + if len(mask_width_ratio_range) != 2: + raise TypeError( + f"mask_width_ratio_range must be a tuple of float and float values: " + f"{mask_width_ratio_range}", + ) + + assert mask_width_ratio_range[1] > mask_width_ratio_range[0] + if isinstance(dim, str): + if dim == "time": + dim = 1 + elif dim == "freq": + dim = 2 + else: + raise ValueError("dim must be int, 'time' or 'freq'") + if dim == 1: + self.mask_axis = "time" + elif dim == 2: + self.mask_axis = "freq" + else: + self.mask_axis = "unknown" + + super().__init__() + self.mask_width_ratio_range = mask_width_ratio_range + self.num_mask = num_mask + self.dim = dim + self.replace_with_zero = replace_with_zero + + def extra_repr(self): + return ( + f"mask_width_ratio_range={self.mask_width_ratio_range}, " + f"num_mask={self.num_mask}, axis={self.mask_axis}" + ) + + def forward(self, spec: torch.Tensor, spec_lengths: torch.Tensor = None): + """Forward function. + + Args: + spec: (Batch, Length, Freq) + """ + + max_seq_len = spec.shape[self.dim] + min_mask_width = math.floor(max_seq_len * self.mask_width_ratio_range[0]) + min_mask_width = max([0, min_mask_width]) + max_mask_width = math.floor(max_seq_len * self.mask_width_ratio_range[1]) + max_mask_width = min([max_seq_len, max_mask_width]) + + if max_mask_width > min_mask_width: + return mask_along_axis( + spec, + spec_lengths, + mask_width_range=(min_mask_width, max_mask_width), + dim=self.dim, + num_mask=self.num_mask, + replace_with_zero=self.replace_with_zero, + ) + return spec, spec_lengths + + +class MaskAlongAxisLFR(torch.nn.Module): + def __init__( + self, + mask_width_range: Union[int, Sequence[int]] = (0, 30), + num_mask: int = 2, + dim: Union[int, str] = "time", + replace_with_zero: bool = True, + lfr_rate: int = 1, + ): + if isinstance(mask_width_range, int): + mask_width_range = (0, mask_width_range) + if len(mask_width_range) != 2: + raise TypeError( + f"mask_width_range must be a tuple of int and int values: " f"{mask_width_range}", + ) + + assert mask_width_range[1] > mask_width_range[0] + if isinstance(dim, str): + if dim == "time": + dim = 1 + lfr_rate = 1 + elif dim == "freq": + dim = 2 + else: + raise ValueError("dim must be int, 'time' or 'freq'") + if dim == 1: + self.mask_axis = "time" + lfr_rate = 1 + elif dim == 2: + self.mask_axis = "freq" + else: + self.mask_axis = "unknown" + + super().__init__() + self.mask_width_range = mask_width_range + self.num_mask = num_mask + self.dim = dim + self.replace_with_zero = replace_with_zero + self.lfr_rate = lfr_rate + + def extra_repr(self): + return ( + f"mask_width_range={self.mask_width_range}, " + f"num_mask={self.num_mask}, axis={self.mask_axis}" + ) + + def forward(self, spec: torch.Tensor, spec_lengths: torch.Tensor = None): + """Forward function. + + Args: + spec: (Batch, Length, Freq) + """ + + return mask_along_axis_lfr( + spec, + spec_lengths, + mask_width_range=self.mask_width_range, + dim=self.dim, + num_mask=self.num_mask, + replace_with_zero=self.replace_with_zero, + lfr_rate=self.lfr_rate, + ) diff --git a/funasr/models/specaug/profileaug.py b/funasr/models/specaug/profileaug.py new file mode 100644 index 0000000000000000000000000000000000000000..669d323f812a6ca2884880da2f8859a137eb3f61 --- /dev/null +++ b/funasr/models/specaug/profileaug.py @@ -0,0 +1,142 @@ +from typing import Tuple, Optional +import numpy as np +import torch +from torch.nn import functional as F +import torch.nn as nn + + +class ProfileAug(nn.Module): + """ + Implement the augmentation for profiles including: + - Split aug: split one profile into two profiles, i.e., main and inaccurate, labels assigned to main + - Merge aug: merge two profiles into one, labels are also merged into one, the other set to zero + - Disturb aug: disturb some profile with others to simulate the inaccurate clustering centroids. + """ + + def __init__( + self, + apply_split_aug: bool = True, + split_aug_prob: float = 0.05, + apply_merge_aug: bool = True, + merge_aug_prob: float = 0.2, + apply_disturb_aug: bool = True, + disturb_aug_prob: float = 0.4, + disturb_alpha: float = 0.2, + ) -> None: + super().__init__() + self.apply_split_aug = apply_split_aug + self.split_aug_prob = split_aug_prob + self.apply_merge_aug = apply_merge_aug + self.merge_aug_prob = merge_aug_prob + self.apply_disturb_aug = apply_disturb_aug + self.disturb_aug_prob = disturb_aug_prob + self.disturb_alpha = disturb_alpha + + def split_aug(self, profile: torch.Tensor, binary_labels: torch.Tensor, mask: torch.Tensor): + # B, N + bsz, dim = profile.shape[0], profile.shape[-1] + profile_norm = torch.linalg.norm(profile, dim=-1, keepdim=False) + spk_count = binary_labels.sum(dim=1) + prob = np.random.rand(bsz) + batch_indices = np.nonzero(prob < self.split_aug_prob)[0] + for idx in batch_indices: + valid_spk_idx = torch.nonzero(spk_count[idx] * mask[idx]) + pad_spk_idx = torch.nonzero((spk_count[idx] == 0) * mask[idx]) + if len(valid_spk_idx) == 0 or len(pad_spk_idx) == 0: + continue + split_spk_idx = valid_spk_idx[torch.randint(len(valid_spk_idx), ())] + to_cover_idx = pad_spk_idx[torch.randint(len(pad_spk_idx), ())] + disturb_vec = torch.randn((dim,)).to(profile) + disturb_vec = F.normalize(disturb_vec, dim=-1) + profile[idx, to_cover_idx] = F.normalize( + profile[idx, split_spk_idx] + self.disturb_alpha * disturb_vec + ) + mask[idx, split_spk_idx] = 0 + mask[idx, to_cover_idx] = 0 + return profile, binary_labels, mask + + def merge_aug(self, profile: torch.Tensor, binary_labels: torch.Tensor, mask: torch.Tensor): + bsz, dim = profile.shape[0], profile.shape[-1] + profile_norm = torch.linalg.norm(profile, dim=-1, keepdim=False) + spk_count = binary_labels.sum(dim=1) + prob = np.random.rand(bsz) + batch_indices = np.nonzero(prob < self.merge_aug_prob)[0] + for idx in batch_indices: + valid_spk_idx = torch.nonzero(profile_norm[idx] * mask[idx]) + if len(valid_spk_idx) == 0: + continue + to_merge = torch.randint(len(valid_spk_idx), (2,)) + spk_idx_1, spk_idx_2 = valid_spk_idx[to_merge[0]], valid_spk_idx[to_merge[1]] + # merge profile + profile[idx, spk_idx_1] = profile[idx, spk_idx_1] + profile[idx, spk_idx_2] + profile[idx, spk_idx_1] = F.normalize(profile[idx, spk_idx_1], dim=-1) + profile[idx, spk_idx_2] = 0 + # merge binary labels + binary_labels[idx, :, spk_idx_1] = ( + binary_labels[idx, :, spk_idx_1] + binary_labels[idx, :, spk_idx_2] + ) + binary_labels[idx, :, spk_idx_1] = (binary_labels[idx, :, spk_idx_1] > 0).to( + binary_labels + ) + binary_labels[idx, :, spk_idx_2] = 0 + + mask[idx, spk_idx_1] = 0 + mask[idx, spk_idx_2] = 0 + + return profile, binary_labels, mask + + def disturb_aug(self, profile: torch.Tensor, binary_labels: torch.Tensor, mask: torch.Tensor): + bsz, dim = profile.shape[0], profile.shape[-1] + profile_norm = torch.linalg.norm(profile, dim=-1, keepdim=False) + spk_count = binary_labels.sum(dim=1) + prob = np.random.rand(bsz) + batch_indices = np.nonzero(prob < self.disturb_aug_prob)[0] + for idx in batch_indices: + pos_spk_idx = torch.nonzero(spk_count[idx] * mask[idx]) + valid_spk_idx = torch.nonzero(profile_norm[idx] * mask[idx]) + if len(pos_spk_idx) == 0 or len(valid_spk_idx) == 0: + continue + to_disturb_idx = pos_spk_idx[torch.randint(len(pos_spk_idx), ())] + disturb_idx = valid_spk_idx[torch.randint(len(valid_spk_idx), ())] + alpha = self.disturb_alpha * torch.rand(()).item() + profile[idx, to_disturb_idx] = (1 - alpha) * profile[ + idx, to_disturb_idx + ] + alpha * profile[idx, disturb_idx] + profile[idx, to_disturb_idx] = F.normalize(profile[idx, to_disturb_idx], dim=-1) + mask[idx, to_disturb_idx] = 0 + + return profile, binary_labels, mask + + def forward( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor = None, + profile: torch.Tensor = None, + profile_lengths: torch.Tensor = None, + binary_labels: torch.Tensor = None, + labels_length: torch.Tensor = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: + + # copy inputs to avoid inplace-operation + speech, profile, binary_labels = ( + torch.clone(speech), + torch.clone(profile), + torch.clone(binary_labels), + ) + profile = F.normalize(profile, dim=-1) + + profile_mask = torch.ones(profile.shape[:2]).to(profile) + if self.apply_disturb_aug: + profile, binary_labels, profile_mask = self.disturb_aug( + profile, binary_labels, profile_mask + ) + if self.apply_split_aug: + profile, binary_labels, profile_mask = self.split_aug( + profile, binary_labels, profile_mask + ) + if self.apply_merge_aug: + profile, binary_labels, profile_mask = self.merge_aug( + profile, binary_labels, profile_mask + ) + + return speech, profile, binary_labels diff --git a/funasr/models/specaug/specaug.py b/funasr/models/specaug/specaug.py new file mode 100644 index 0000000000000000000000000000000000000000..75a4b0b28698d679f4446242c0e7756c9b5bcfb5 --- /dev/null +++ b/funasr/models/specaug/specaug.py @@ -0,0 +1,186 @@ +"""SpecAugment module.""" + +from typing import Optional +from typing import Sequence +from typing import Union + +from funasr.models.specaug.mask_along_axis import MaskAlongAxis +from funasr.models.specaug.mask_along_axis import MaskAlongAxisVariableMaxWidth +from funasr.models.specaug.mask_along_axis import MaskAlongAxisLFR +from funasr.models.specaug.time_warp import TimeWarp +from funasr.register import tables + +import torch.nn as nn + + +@tables.register("specaug_classes", "SpecAug") +class SpecAug(nn.Module): + """Implementation of SpecAug. + + Reference: + Daniel S. Park et al. + "SpecAugment: A Simple Data + Augmentation Method for Automatic Speech Recognition" + + .. warning:: + When using cuda mode, time_warp doesn't have reproducibility + due to `torch.nn.functional.interpolate`. + + """ + + def __init__( + self, + apply_time_warp: bool = True, + time_warp_window: int = 5, + time_warp_mode: str = "bicubic", + apply_freq_mask: bool = True, + freq_mask_width_range: Union[int, Sequence[int]] = (0, 20), + num_freq_mask: int = 2, + apply_time_mask: bool = True, + time_mask_width_range: Optional[Union[int, Sequence[int]]] = None, + time_mask_width_ratio_range: Optional[Union[float, Sequence[float]]] = None, + num_time_mask: int = 2, + ): + if not apply_time_warp and not apply_time_mask and not apply_freq_mask: + raise ValueError("Either one of time_warp, time_mask, or freq_mask should be applied") + if ( + apply_time_mask + and (time_mask_width_range is not None) + and (time_mask_width_ratio_range is not None) + ): + raise ValueError( + 'Either one of "time_mask_width_range" or ' + '"time_mask_width_ratio_range" can be used' + ) + super().__init__() + self.apply_time_warp = apply_time_warp + self.apply_freq_mask = apply_freq_mask + self.apply_time_mask = apply_time_mask + + if apply_time_warp: + self.time_warp = TimeWarp(window=time_warp_window, mode=time_warp_mode) + else: + self.time_warp = None + + if apply_freq_mask: + self.freq_mask = MaskAlongAxis( + dim="freq", + mask_width_range=freq_mask_width_range, + num_mask=num_freq_mask, + ) + else: + self.freq_mask = None + + if apply_time_mask: + if time_mask_width_range is not None: + self.time_mask = MaskAlongAxis( + dim="time", + mask_width_range=time_mask_width_range, + num_mask=num_time_mask, + ) + elif time_mask_width_ratio_range is not None: + self.time_mask = MaskAlongAxisVariableMaxWidth( + dim="time", + mask_width_ratio_range=time_mask_width_ratio_range, + num_mask=num_time_mask, + ) + else: + raise ValueError( + 'Either one of "time_mask_width_range" or ' + '"time_mask_width_ratio_range" should be used.' + ) + else: + self.time_mask = None + + def forward(self, x, x_lengths=None): + if self.time_warp is not None: + x, x_lengths = self.time_warp(x, x_lengths) + if self.freq_mask is not None: + x, x_lengths = self.freq_mask(x, x_lengths) + if self.time_mask is not None: + x, x_lengths = self.time_mask(x, x_lengths) + return x, x_lengths + + +@tables.register("specaug_classes", "SpecAugLFR") +class SpecAugLFR(nn.Module): + """Implementation of SpecAug. + lfr_rate：low frame rate + """ + + def __init__( + self, + apply_time_warp: bool = True, + time_warp_window: int = 5, + time_warp_mode: str = "bicubic", + apply_freq_mask: bool = True, + freq_mask_width_range: Union[int, Sequence[int]] = (0, 20), + num_freq_mask: int = 2, + lfr_rate: int = 0, + apply_time_mask: bool = True, + time_mask_width_range: Optional[Union[int, Sequence[int]]] = None, + time_mask_width_ratio_range: Optional[Union[float, Sequence[float]]] = None, + num_time_mask: int = 2, + ): + if not apply_time_warp and not apply_time_mask and not apply_freq_mask: + raise ValueError("Either one of time_warp, time_mask, or freq_mask should be applied") + if ( + apply_time_mask + and (time_mask_width_range is not None) + and (time_mask_width_ratio_range is not None) + ): + raise ValueError( + 'Either one of "time_mask_width_range" or ' + '"time_mask_width_ratio_range" can be used' + ) + super().__init__() + self.apply_time_warp = apply_time_warp + self.apply_freq_mask = apply_freq_mask + self.apply_time_mask = apply_time_mask + + if apply_time_warp: + self.time_warp = TimeWarp(window=time_warp_window, mode=time_warp_mode) + else: + self.time_warp = None + + if apply_freq_mask: + self.freq_mask = MaskAlongAxisLFR( + dim="freq", + mask_width_range=freq_mask_width_range, + num_mask=num_freq_mask, + lfr_rate=lfr_rate + 1, + ) + + else: + self.freq_mask = None + + if apply_time_mask: + if time_mask_width_range is not None: + self.time_mask = MaskAlongAxisLFR( + dim="time", + mask_width_range=time_mask_width_range, + num_mask=num_time_mask, + lfr_rate=lfr_rate + 1, + ) + elif time_mask_width_ratio_range is not None: + self.time_mask = MaskAlongAxisVariableMaxWidth( + dim="time", + mask_width_ratio_range=time_mask_width_ratio_range, + num_mask=num_time_mask, + ) + else: + raise ValueError( + 'Either one of "time_mask_width_range" or ' + '"time_mask_width_ratio_range" should be used.' + ) + else: + self.time_mask = None + + def forward(self, x, x_lengths=None): + if self.time_warp is not None: + x, x_lengths = self.time_warp(x, x_lengths) + if self.freq_mask is not None: + x, x_lengths = self.freq_mask(x, x_lengths) + if self.time_mask is not None: + x, x_lengths = self.time_mask(x, x_lengths) + return x, x_lengths diff --git a/funasr/models/specaug/time_warp.py b/funasr/models/specaug/time_warp.py new file mode 100644 index 0000000000000000000000000000000000000000..ef3532f6af6cef897490fddefe3c6650f3af6d3f --- /dev/null +++ b/funasr/models/specaug/time_warp.py @@ -0,0 +1,89 @@ +"""Time warp module.""" + +import torch + +from funasr.models.transformer.utils.nets_utils import pad_list + +DEFAULT_TIME_WARP_MODE = "bicubic" + + +def time_warp(x: torch.Tensor, window: int = 80, mode: str = DEFAULT_TIME_WARP_MODE): + """Time warping using torch.interpolate. + + Args: + x: (Batch, Time, Freq) + window: time warp parameter + mode: Interpolate mode + """ + + # bicubic supports 4D or more dimension tensor + org_size = x.size() + if x.dim() == 3: + # x: (Batch, Time, Freq) -> (Batch, 1, Time, Freq) + x = x[:, None] + + t = x.shape[2] + if t - window <= window: + return x.view(*org_size) + + center = torch.randint(window, t - window, (1,))[0] + warped = torch.randint(center - window, center + window, (1,))[0] + 1 + + # left: (Batch, Channel, warped, Freq) + # right: (Batch, Channel, time - warped, Freq) + left = torch.nn.functional.interpolate( + x[:, :, :center], (warped, x.shape[3]), mode=mode, align_corners=False + ) + right = torch.nn.functional.interpolate( + x[:, :, center:], (t - warped, x.shape[3]), mode=mode, align_corners=False + ) + + if x.requires_grad: + x = torch.cat([left, right], dim=-2) + else: + x[:, :, :warped] = left + x[:, :, warped:] = right + + return x.view(*org_size) + + +class TimeWarp(torch.nn.Module): + """Time warping using torch.interpolate. + + Args: + window: time warp parameter + mode: Interpolate mode + """ + + def __init__(self, window: int = 80, mode: str = DEFAULT_TIME_WARP_MODE): + super().__init__() + self.window = window + self.mode = mode + + def extra_repr(self): + return f"window={self.window}, mode={self.mode}" + + def forward(self, x: torch.Tensor, x_lengths: torch.Tensor = None): + """Forward function. + + Args: + x: (Batch, Time, Freq) + x_lengths: (Batch,) + """ + + if x_lengths is None or all(le == x_lengths[0] for le in x_lengths): + # Note that applying same warping for each sample + y = time_warp(x, window=self.window, mode=self.mode) + else: + # FIXME(kamo): I have no idea to batchify Timewarp + ys = [] + for i in range(x.size(0)): + _y = time_warp( + x[i][None, : x_lengths[i]], + window=self.window, + mode=self.mode, + )[0] + ys.append(_y) + y = pad_list(ys, 0.0) + + return y, x_lengths diff --git a/funasr/models/transducer/__init__.py b/funasr/models/transducer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/models/transducer/beam_search_transducer.py b/funasr/models/transducer/beam_search_transducer.py new file mode 100644 index 0000000000000000000000000000000000000000..aa9e34e1e4159446e7e6073c3b86875ef4bd625c --- /dev/null +++ b/funasr/models/transducer/beam_search_transducer.py @@ -0,0 +1,694 @@ +#!/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 numpy as np +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple, Union + +from funasr.models.transducer.joint_network import JointNetwork + + +@dataclass +class Hypothesis: + """Default hypothesis definition for Transducer search algorithms. + + Args: + score: Total log-probability. + yseq: Label sequence as integer ID sequence. + dec_state: RNNDecoder or StatelessDecoder state. + ((N, 1, D_dec), (N, 1, D_dec) or None) or None + lm_state: RNNLM state. ((N, D_lm), (N, D_lm)) or None + + """ + + score: float + yseq: List[int] + dec_state: Optional[Tuple[torch.Tensor, Optional[torch.Tensor]]] = None + lm_state: Optional[Union[Dict[str, Any], List[Any]]] = None + + +@dataclass +class ExtendedHypothesis(Hypothesis): + """Extended hypothesis definition for NSC beam search and mAES. + + Args: + : Hypothesis dataclass arguments. + dec_out: Decoder output sequence. (B, D_dec) + lm_score: Log-probabilities of the LM for given label. (vocab_size) + + """ + + dec_out: torch.Tensor = None + lm_score: torch.Tensor = None + + +class BeamSearchTransducer: + """Beam search implementation for Transducer. + + Args: + decoder: Decoder module. + joint_network: Joint network module. + beam_size: Size of the beam. + lm: LM class. + lm_weight: LM weight for soft fusion. + search_type: Search algorithm to use during inference. + max_sym_exp: Number of maximum symbol expansions at each time step. (TSD) + u_max: Maximum expected target sequence length. (ALSD) + nstep: Number of maximum expansion steps at each time step. (mAES) + expansion_gamma: Allowed logp difference for prune-by-value method. (mAES) + expansion_beta: + Number of additional candidates for expanded hypotheses selection. (mAES) + score_norm: Normalize final scores by length. + nbest: Number of final hypothesis. + streaming: Whether to perform chunk-by-chunk beam search. + + """ + + def __init__( + self, + decoder, + joint_network: JointNetwork, + beam_size: int, + lm: Optional[torch.nn.Module] = None, + lm_weight: float = 0.1, + search_type: str = "default", + max_sym_exp: int = 3, + u_max: int = 50, + nstep: int = 2, + expansion_gamma: float = 2.3, + expansion_beta: int = 2, + score_norm: bool = False, + nbest: int = 1, + streaming: bool = False, + ) -> None: + """Construct a BeamSearchTransducer object.""" + super().__init__() + + self.decoder = decoder + self.joint_network = joint_network + + self.vocab_size = decoder.vocab_size + + assert ( + beam_size <= self.vocab_size + ), "beam_size (%d) should be smaller than or equal to vocabulary size (%d)." % ( + beam_size, + self.vocab_size, + ) + self.beam_size = beam_size + + if search_type == "default": + self.search_algorithm = self.default_beam_search + elif search_type == "tsd": + assert max_sym_exp > 1, "max_sym_exp (%d) should be greater than one." % (max_sym_exp) + self.max_sym_exp = max_sym_exp + + self.search_algorithm = self.time_sync_decoding + elif search_type == "alsd": + assert not streaming, "ALSD is not available in streaming mode." + + assert u_max >= 0, "u_max should be a positive integer, a portion of max_T." + self.u_max = u_max + + self.search_algorithm = self.align_length_sync_decoding + elif search_type == "maes": + assert self.vocab_size >= beam_size + expansion_beta, ( + "beam_size (%d) + expansion_beta (%d) " + " should be smaller than or equal to vocab size (%d)." + % (beam_size, expansion_beta, self.vocab_size) + ) + self.max_candidates = beam_size + expansion_beta + + self.nstep = nstep + self.expansion_gamma = expansion_gamma + + self.search_algorithm = self.modified_adaptive_expansion_search + else: + raise NotImplementedError("Specified search type (%s) is not supported." % search_type) + + self.use_lm = lm is not None + + if self.use_lm: + assert hasattr(lm, "rnn_type"), "Transformer LM is currently not supported." + + self.sos = self.vocab_size - 1 + + self.lm = lm + self.lm_weight = lm_weight + + self.score_norm = score_norm + self.nbest = nbest + + self.reset_inference_cache() + + def __call__( + self, + enc_out: torch.Tensor, + is_final: bool = True, + ) -> List[Hypothesis]: + """Perform beam search. + + Args: + enc_out: Encoder output sequence. (T, D_enc) + is_final: Whether enc_out is the final chunk of data. + + Returns: + nbest_hyps: N-best decoding results + + """ + self.decoder.set_device(enc_out.device) + + hyps = self.search_algorithm(enc_out) + + if is_final: + self.reset_inference_cache() + + return self.sort_nbest(hyps) + + self.search_cache = hyps + + return hyps + + def reset_inference_cache(self) -> None: + """Reset cache for decoder scoring and streaming.""" + self.decoder.score_cache = {} + self.search_cache = None + + def sort_nbest(self, hyps: List[Hypothesis]) -> List[Hypothesis]: + """Sort in-place hypotheses by score or score given sequence length. + + Args: + hyps: Hypothesis. + + Return: + hyps: Sorted hypothesis. + + """ + if self.score_norm: + hyps.sort(key=lambda x: x.score / len(x.yseq), reverse=True) + else: + hyps.sort(key=lambda x: x.score, reverse=True) + + return hyps[: self.nbest] + + def recombine_hyps(self, hyps: List[Hypothesis]) -> List[Hypothesis]: + """Recombine hypotheses with same label ID sequence. + + Args: + hyps: Hypotheses. + + Returns: + final: Recombined hypotheses. + + """ + final = {} + + for hyp in hyps: + str_yseq = "_".join(map(str, hyp.yseq)) + + if str_yseq in final: + final[str_yseq].score = np.logaddexp(final[str_yseq].score, hyp.score) + else: + final[str_yseq] = hyp + + return [*final.values()] + + def select_k_expansions( + self, + hyps: List[ExtendedHypothesis], + topk_idx: torch.Tensor, + topk_logp: torch.Tensor, + ) -> List[ExtendedHypothesis]: + """Return K hypotheses candidates for expansion from a list of hypothesis. + + K candidates are selected according to the extended hypotheses probabilities + and a prune-by-value method. Where K is equal to beam_size + beta. + + Args: + hyps: Hypotheses. + topk_idx: Indices of candidates hypothesis. + topk_logp: Log-probabilities of candidates hypothesis. + + Returns: + k_expansions: Best K expansion hypotheses candidates. + + """ + k_expansions = [] + + for i, hyp in enumerate(hyps): + hyp_i = [(int(k), hyp.score + float(v)) for k, v in zip(topk_idx[i], topk_logp[i])] + k_best_exp = max(hyp_i, key=lambda x: x[1])[1] + + k_expansions.append( + sorted( + filter(lambda x: (k_best_exp - self.expansion_gamma) <= x[1], hyp_i), + key=lambda x: x[1], + reverse=True, + ) + ) + + return k_expansions + + def create_lm_batch_inputs(self, hyps_seq: List[List[int]]) -> torch.Tensor: + """Make batch of inputs with left padding for LM scoring. + + Args: + hyps_seq: Hypothesis sequences. + + Returns: + : Padded batch of sequences. + + """ + max_len = max([len(h) for h in hyps_seq]) + + return torch.LongTensor( + [[self.sos] + ([0] * (max_len - len(h))) + h[1:] for h in hyps_seq], + device=self.decoder.device, + ) + + def default_beam_search(self, enc_out: torch.Tensor) -> List[Hypothesis]: + """Beam search implementation without prefix search. + + Modified from https://arxiv.org/pdf/1211.3711.pdf + + Args: + enc_out: Encoder output sequence. (T, D) + + Returns: + nbest_hyps: N-best hypothesis. + + """ + beam_k = min(self.beam_size, (self.vocab_size - 1)) + max_t = len(enc_out) + + if self.search_cache is not None: + kept_hyps = self.search_cache + else: + kept_hyps = [ + Hypothesis( + score=0.0, + yseq=[0], + dec_state=self.decoder.init_state(1), + ) + ] + + for t in range(max_t): + hyps = kept_hyps + kept_hyps = [] + + while True: + max_hyp = max(hyps, key=lambda x: x.score) + hyps.remove(max_hyp) + + label = torch.full( + (1, 1), + max_hyp.yseq[-1], + dtype=torch.long, + device=self.decoder.device, + ) + dec_out, state = self.decoder.score( + label, + max_hyp.yseq, + max_hyp.dec_state, + ) + + logp = torch.log_softmax( + self.joint_network(enc_out[t : t + 1, :], dec_out), + dim=-1, + ).squeeze(0) + top_k = logp[1:].topk(beam_k, dim=-1) + + kept_hyps.append( + Hypothesis( + score=(max_hyp.score + float(logp[0:1])), + yseq=max_hyp.yseq, + dec_state=max_hyp.dec_state, + lm_state=max_hyp.lm_state, + ) + ) + + if self.use_lm: + lm_scores, lm_state = self.lm.score( + torch.LongTensor([self.sos] + max_hyp.yseq[1:], device=self.decoder.device), + max_hyp.lm_state, + None, + ) + else: + lm_state = max_hyp.lm_state + + for logp, k in zip(*top_k): + score = max_hyp.score + float(logp) + + if self.use_lm: + score += self.lm_weight * lm_scores[k + 1] + + hyps.append( + Hypothesis( + score=score, + yseq=max_hyp.yseq + [int(k + 1)], + dec_state=state, + lm_state=lm_state, + ) + ) + + hyps_max = float(max(hyps, key=lambda x: x.score).score) + kept_most_prob = sorted( + [hyp for hyp in kept_hyps if hyp.score > hyps_max], + key=lambda x: x.score, + ) + if len(kept_most_prob) >= self.beam_size: + kept_hyps = kept_most_prob + break + + return kept_hyps + + def align_length_sync_decoding( + self, + enc_out: torch.Tensor, + ) -> List[Hypothesis]: + """Alignment-length synchronous beam search implementation. + + Based on https://ieeexplore.ieee.org/document/9053040 + + Args: + h: Encoder output sequences. (T, D) + + Returns: + nbest_hyps: N-best hypothesis. + + """ + t_max = int(enc_out.size(0)) + u_max = min(self.u_max, (t_max - 1)) + + B = [Hypothesis(yseq=[0], score=0.0, dec_state=self.decoder.init_state(1))] + final = [] + + if self.use_lm: + B[0].lm_state = self.lm.zero_state() + + for i in range(t_max + u_max): + A = [] + + B_ = [] + B_enc_out = [] + for hyp in B: + u = len(hyp.yseq) - 1 + t = i - u + + if t > (t_max - 1): + continue + + B_.append(hyp) + B_enc_out.append((t, enc_out[t])) + + if B_: + beam_enc_out = torch.stack([b[1] for b in B_enc_out]) + beam_dec_out, beam_state = self.decoder.batch_score(B_) + + beam_logp = torch.log_softmax( + self.joint_network(beam_enc_out, beam_dec_out), + dim=-1, + ) + beam_topk = beam_logp[:, 1:].topk(self.beam_size, dim=-1) + + if self.use_lm: + beam_lm_scores, beam_lm_states = self.lm.batch_score( + self.create_lm_batch_inputs([b.yseq for b in B_]), + [b.lm_state for b in B_], + None, + ) + + for i, hyp in enumerate(B_): + new_hyp = Hypothesis( + score=(hyp.score + float(beam_logp[i, 0])), + yseq=hyp.yseq[:], + dec_state=hyp.dec_state, + lm_state=hyp.lm_state, + ) + + A.append(new_hyp) + + if B_enc_out[i][0] == (t_max - 1): + final.append(new_hyp) + + for logp, k in zip(beam_topk[0][i], beam_topk[1][i] + 1): + new_hyp = Hypothesis( + score=(hyp.score + float(logp)), + yseq=(hyp.yseq[:] + [int(k)]), + dec_state=self.decoder.select_state(beam_state, i), + lm_state=hyp.lm_state, + ) + + if self.use_lm: + new_hyp.score += self.lm_weight * beam_lm_scores[i, k] + new_hyp.lm_state = beam_lm_states[i] + + A.append(new_hyp) + + B = sorted(A, key=lambda x: x.score, reverse=True)[: self.beam_size] + B = self.recombine_hyps(B) + + if final: + return final + + return B + + def time_sync_decoding(self, enc_out: torch.Tensor) -> List[Hypothesis]: + """Time synchronous beam search implementation. + + Based on https://ieeexplore.ieee.org/document/9053040 + + Args: + enc_out: Encoder output sequence. (T, D) + + Returns: + nbest_hyps: N-best hypothesis. + + """ + if self.search_cache is not None: + B = self.search_cache + else: + B = [ + Hypothesis( + yseq=[0], + score=0.0, + dec_state=self.decoder.init_state(1), + ) + ] + + if self.use_lm: + B[0].lm_state = self.lm.zero_state() + + for enc_out_t in enc_out: + A = [] + C = B + + enc_out_t = enc_out_t.unsqueeze(0) + + for v in range(self.max_sym_exp): + D = [] + + beam_dec_out, beam_state = self.decoder.batch_score(C) + + beam_logp = torch.log_softmax( + self.joint_network(enc_out_t, beam_dec_out), + dim=-1, + ) + beam_topk = beam_logp[:, 1:].topk(self.beam_size, dim=-1) + + seq_A = [h.yseq for h in A] + + for i, hyp in enumerate(C): + if hyp.yseq not in seq_A: + A.append( + Hypothesis( + score=(hyp.score + float(beam_logp[i, 0])), + yseq=hyp.yseq[:], + dec_state=hyp.dec_state, + lm_state=hyp.lm_state, + ) + ) + else: + dict_pos = seq_A.index(hyp.yseq) + + A[dict_pos].score = np.logaddexp( + A[dict_pos].score, (hyp.score + float(beam_logp[i, 0])) + ) + + if v < (self.max_sym_exp - 1): + if self.use_lm: + beam_lm_scores, beam_lm_states = self.lm.batch_score( + self.create_lm_batch_inputs([c.yseq for c in C]), + [c.lm_state for c in C], + None, + ) + + for i, hyp in enumerate(C): + for logp, k in zip(beam_topk[0][i], beam_topk[1][i] + 1): + new_hyp = Hypothesis( + score=(hyp.score + float(logp)), + yseq=(hyp.yseq + [int(k)]), + dec_state=self.decoder.select_state(beam_state, i), + lm_state=hyp.lm_state, + ) + + if self.use_lm: + new_hyp.score += self.lm_weight * beam_lm_scores[i, k] + new_hyp.lm_state = beam_lm_states[i] + + D.append(new_hyp) + + C = sorted(D, key=lambda x: x.score, reverse=True)[: self.beam_size] + + B = sorted(A, key=lambda x: x.score, reverse=True)[: self.beam_size] + + return B + + def modified_adaptive_expansion_search( + self, + enc_out: torch.Tensor, + ) -> List[ExtendedHypothesis]: + """Modified version of Adaptive Expansion Search (mAES). + + Based on AES (https://ieeexplore.ieee.org/document/9250505) and + NSC (https://arxiv.org/abs/2201.05420). + + Args: + enc_out: Encoder output sequence. (T, D_enc) + + Returns: + nbest_hyps: N-best hypothesis. + + """ + if self.search_cache is not None: + kept_hyps = self.search_cache + else: + init_tokens = [ + ExtendedHypothesis( + yseq=[0], + score=0.0, + dec_state=self.decoder.init_state(1), + ) + ] + + beam_dec_out, beam_state = self.decoder.batch_score( + init_tokens, + ) + + if self.use_lm: + beam_lm_scores, beam_lm_states = self.lm.batch_score( + self.create_lm_batch_inputs([h.yseq for h in init_tokens]), + [h.lm_state for h in init_tokens], + None, + ) + + lm_state = beam_lm_states[0] + lm_score = beam_lm_scores[0] + else: + lm_state = None + lm_score = None + + kept_hyps = [ + ExtendedHypothesis( + yseq=[0], + score=0.0, + dec_state=self.decoder.select_state(beam_state, 0), + dec_out=beam_dec_out[0], + lm_state=lm_state, + lm_score=lm_score, + ) + ] + + for enc_out_t in enc_out: + hyps = kept_hyps + kept_hyps = [] + + beam_enc_out = enc_out_t.unsqueeze(0) + + list_b = [] + for n in range(self.nstep): + beam_dec_out = torch.stack([h.dec_out for h in hyps]) + + beam_logp, beam_idx = torch.log_softmax( + self.joint_network(beam_enc_out, beam_dec_out), + dim=-1, + ).topk(self.max_candidates, dim=-1) + + k_expansions = self.select_k_expansions(hyps, beam_idx, beam_logp) + + list_exp = [] + for i, hyp in enumerate(hyps): + for k, new_score in k_expansions[i]: + new_hyp = ExtendedHypothesis( + yseq=hyp.yseq[:], + score=new_score, + dec_out=hyp.dec_out, + dec_state=hyp.dec_state, + lm_state=hyp.lm_state, + lm_score=hyp.lm_score, + ) + + if k == 0: + list_b.append(new_hyp) + else: + new_hyp.yseq.append(int(k)) + + if self.use_lm: + new_hyp.score += self.lm_weight * float(hyp.lm_score[k]) + + list_exp.append(new_hyp) + + if not list_exp: + kept_hyps = sorted( + self.recombine_hyps(list_b), key=lambda x: x.score, reverse=True + )[: self.beam_size] + + break + else: + beam_dec_out, beam_state = self.decoder.batch_score( + list_exp, + ) + + if self.use_lm: + beam_lm_scores, beam_lm_states = self.lm.batch_score( + self.create_lm_batch_inputs([h.yseq for h in list_exp]), + [h.lm_state for h in list_exp], + None, + ) + + if n < (self.nstep - 1): + for i, hyp in enumerate(list_exp): + hyp.dec_out = beam_dec_out[i] + hyp.dec_state = self.decoder.select_state(beam_state, i) + + if self.use_lm: + hyp.lm_state = beam_lm_states[i] + hyp.lm_score = beam_lm_scores[i] + + hyps = list_exp[:] + else: + beam_logp = torch.log_softmax( + self.joint_network(beam_enc_out, beam_dec_out), + dim=-1, + ) + + for i, hyp in enumerate(list_exp): + hyp.score += float(beam_logp[i, 0]) + + hyp.dec_out = beam_dec_out[i] + hyp.dec_state = self.decoder.select_state(beam_state, i) + + if self.use_lm: + hyp.lm_state = beam_lm_states[i] + hyp.lm_score = beam_lm_scores[i] + + kept_hyps = sorted( + self.recombine_hyps(list_b + list_exp), + key=lambda x: x.score, + reverse=True, + )[: self.beam_size] + + return kept_hyps diff --git a/funasr/models/transducer/joint_network.py b/funasr/models/transducer/joint_network.py new file mode 100644 index 0000000000000000000000000000000000000000..955d6a0cfdafdaa07b84ddc5d19ef178010ae5af --- /dev/null +++ b/funasr/models/transducer/joint_network.py @@ -0,0 +1,64 @@ +#!/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 get_activation + + +@tables.register("joint_network_classes", "joint_network") +class JointNetwork(torch.nn.Module): + """Transducer joint network module. + + Args: + output_size: Output size. + encoder_size: Encoder output size. + decoder_size: Decoder output size.. + joint_space_size: Joint space size. + joint_act_type: Type of activation for joint network. + **activation_parameters: Parameters for the activation function. + + """ + + def __init__( + self, + output_size: int, + encoder_size: int, + decoder_size: int, + joint_space_size: int = 256, + joint_activation_type: str = "tanh", + ) -> None: + """Construct a JointNetwork object.""" + super().__init__() + + self.lin_enc = torch.nn.Linear(encoder_size, joint_space_size) + self.lin_dec = torch.nn.Linear(decoder_size, joint_space_size, bias=False) + + self.lin_out = torch.nn.Linear(joint_space_size, output_size) + + self.joint_activation = get_activation(joint_activation_type) + + def forward( + self, + enc_out: torch.Tensor, + dec_out: torch.Tensor, + project_input: bool = True, + ) -> torch.Tensor: + """Joint computation of encoder and decoder hidden state sequences. + + Args: + enc_out: Expanded encoder output state sequences (B, T, 1, D_enc) + dec_out: Expanded decoder output state sequences (B, 1, U, D_dec) + + Returns: + joint_out: Joint output state sequences. (B, T, U, D_out) + + """ + if project_input: + joint_out = self.joint_activation(self.lin_enc(enc_out) + self.lin_dec(dec_out)) + else: + joint_out = self.joint_activation(enc_out + dec_out) + return self.lin_out(joint_out) diff --git a/funasr/models/transducer/model.py b/funasr/models/transducer/model.py new file mode 100644 index 0000000000000000000000000000000000000000..37925fcf661b4f7867b849ffa4cdb10f13fe26cf --- /dev/null +++ b/funasr/models/transducer/model.py @@ -0,0 +1,536 @@ +#!/usr/bin/env python3 +# -*- encoding: utf-8 -*- +# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved. +# MIT License (https://opensource.org/licenses/MIT) + +import time +import torch +import logging +from contextlib import contextmanager +from typing import Dict, Optional, Tuple +from distutils.version import LooseVersion + +from funasr.register import tables +from funasr.utils import postprocess_utils +from funasr.utils.datadir_writer import DatadirWriter +from funasr.train_utils.device_funcs import force_gatherable +from funasr.models.transformer.scorers.ctc import CTCPrefixScorer +from funasr.losses.label_smoothing_loss import LabelSmoothingLoss +from funasr.models.transformer.scorers.length_bonus import LengthBonus +from funasr.models.transformer.utils.nets_utils import get_transducer_task_io +from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank +from funasr.models.transducer.beam_search_transducer import BeamSearchTransducer + + +if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): + from torch.cuda.amp import autocast +else: + # Nothing to do if torch<1.6.0 + @contextmanager + def autocast(enabled=True): + yield + + +@tables.register("model_classes", "Transducer") +class Transducer(torch.nn.Module): + def __init__( + self, + frontend: Optional[str] = None, + frontend_conf: Optional[Dict] = None, + specaug: Optional[str] = None, + specaug_conf: Optional[Dict] = None, + normalize: str = None, + normalize_conf: Optional[Dict] = None, + encoder: str = None, + encoder_conf: Optional[Dict] = None, + decoder: str = None, + decoder_conf: Optional[Dict] = None, + joint_network: str = None, + joint_network_conf: Optional[Dict] = None, + transducer_weight: float = 1.0, + fastemit_lambda: float = 0.0, + auxiliary_ctc_weight: float = 0.0, + auxiliary_ctc_dropout_rate: float = 0.0, + auxiliary_lm_loss_weight: float = 0.0, + auxiliary_lm_loss_smoothing: float = 0.0, + input_size: int = 80, + vocab_size: int = -1, + ignore_id: int = -1, + blank_id: int = 0, + sos: int = 1, + eos: int = 2, + lsm_weight: float = 0.0, + length_normalized_loss: bool = False, + # report_cer: bool = True, + # report_wer: bool = True, + # sym_space: str = "", + # sym_blank: str = "", + # extract_feats_in_collect_stats: bool = True, + share_embedding: bool = False, + # preencoder: Optional[AbsPreEncoder] = None, + # postencoder: Optional[AbsPostEncoder] = None, + **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, + **decoder_conf, + ) + decoder_output_size = decoder.output_size + + joint_network_class = tables.joint_network_classes.get(joint_network) + joint_network = joint_network_class( + vocab_size, + encoder_output_size, + decoder_output_size, + **joint_network_conf, + ) + + self.criterion_transducer = None + self.error_calculator = None + + self.use_auxiliary_ctc = auxiliary_ctc_weight > 0 + self.use_auxiliary_lm_loss = auxiliary_lm_loss_weight > 0 + + if self.use_auxiliary_ctc: + self.ctc_lin = torch.nn.Linear(encoder.output_size(), vocab_size) + self.ctc_dropout_rate = auxiliary_ctc_dropout_rate + + if self.use_auxiliary_lm_loss: + self.lm_lin = torch.nn.Linear(decoder.output_size, vocab_size) + self.lm_loss_smoothing = auxiliary_lm_loss_smoothing + + self.transducer_weight = transducer_weight + self.fastemit_lambda = fastemit_lambda + + self.auxiliary_ctc_weight = auxiliary_ctc_weight + self.auxiliary_lm_loss_weight = auxiliary_lm_loss_weight + self.blank_id = blank_id + self.sos = sos if sos is not None else vocab_size - 1 + self.eos = eos if eos is not None else vocab_size - 1 + self.vocab_size = vocab_size + self.ignore_id = ignore_id + self.frontend = frontend + self.specaug = specaug + self.normalize = normalize + self.encoder = encoder + self.decoder = decoder + self.joint_network = joint_network + + self.criterion_att = LabelSmoothingLoss( + size=vocab_size, + padding_idx=ignore_id, + smoothing=lsm_weight, + normalize_length=length_normalized_loss, + ) + + self.length_normalized_loss = length_normalized_loss + self.beam_search = None + self.ctc = None + self.ctc_weight = 0.0 + + def forward( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + text: torch.Tensor, + text_lengths: torch.Tensor, + **kwargs, + ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]: + """Encoder + Decoder + Calc loss + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + text: (Batch, Length) + text_lengths: (Batch,) + """ + if len(text_lengths.size()) > 1: + text_lengths = text_lengths[:, 0] + if len(speech_lengths.size()) > 1: + speech_lengths = speech_lengths[:, 0] + + batch_size = speech.shape[0] + # 1. Encoder + encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) + if ( + hasattr(self.encoder, "overlap_chunk_cls") + and 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 + ) + # 2. Transducer-related I/O preparation + decoder_in, target, t_len, u_len = get_transducer_task_io( + text, + encoder_out_lens, + ignore_id=self.ignore_id, + ) + + # 3. Decoder + self.decoder.set_device(encoder_out.device) + decoder_out = self.decoder(decoder_in, u_len) + + # 4. Joint Network + joint_out = self.joint_network(encoder_out.unsqueeze(2), decoder_out.unsqueeze(1)) + + # 5. Losses + loss_trans, cer_trans, wer_trans = self._calc_transducer_loss( + encoder_out, + joint_out, + target, + t_len, + u_len, + ) + + loss_ctc, loss_lm = 0.0, 0.0 + + if self.use_auxiliary_ctc: + loss_ctc = self._calc_ctc_loss( + encoder_out, + target, + t_len, + u_len, + ) + + if self.use_auxiliary_lm_loss: + loss_lm = self._calc_lm_loss(decoder_out, target) + + loss = ( + self.transducer_weight * loss_trans + + self.auxiliary_ctc_weight * loss_ctc + + self.auxiliary_lm_loss_weight * loss_lm + ) + + stats = dict( + loss=loss.detach(), + loss_transducer=loss_trans.detach(), + aux_ctc_loss=loss_ctc.detach() if loss_ctc > 0.0 else None, + aux_lm_loss=loss_lm.detach() if loss_lm > 0.0 else None, + cer_transducer=cer_trans, + wer_transducer=wer_trans, + ) + + # force_gatherable: to-device and to-tensor if scalar for DataParallel + loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) + + return loss, stats, weight + + def encode( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Frontend + Encoder. Note that this method is used by asr_inference.py + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + ind: int + """ + with autocast(False): + + # Data augmentation + if self.specaug is not None and self.training: + speech, speech_lengths = self.specaug(speech, speech_lengths) + + # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN + if self.normalize is not None: + speech, speech_lengths = self.normalize(speech, speech_lengths) + + # Forward encoder + # feats: (Batch, Length, Dim) + # -> encoder_out: (Batch, Length2, Dim2) + encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths) + intermediate_outs = None + if isinstance(encoder_out, tuple): + intermediate_outs = encoder_out[1] + encoder_out = encoder_out[0] + + if intermediate_outs is not None: + return (encoder_out, intermediate_outs), encoder_out_lens + + return encoder_out, encoder_out_lens + + def _calc_transducer_loss( + self, + encoder_out: torch.Tensor, + joint_out: torch.Tensor, + target: torch.Tensor, + t_len: torch.Tensor, + u_len: torch.Tensor, + ) -> Tuple[torch.Tensor, Optional[float], Optional[float]]: + """Compute Transducer loss. + + Args: + encoder_out: Encoder output sequences. (B, T, D_enc) + joint_out: Joint Network output sequences (B, T, U, D_joint) + target: Target label ID sequences. (B, L) + t_len: Encoder output sequences lengths. (B,) + u_len: Target label ID sequences lengths. (B,) + + Return: + loss_transducer: Transducer loss value. + cer_transducer: Character error rate for Transducer. + wer_transducer: Word Error Rate for Transducer. + + """ + if self.criterion_transducer is None: + try: + from warp_rnnt import rnnt_loss as RNNTLoss + + self.criterion_transducer = RNNTLoss + + except ImportError: + logging.error( + "warp-rnnt was not installed." "Please consult the installation documentation." + ) + exit(1) + + log_probs = torch.log_softmax(joint_out, dim=-1) + + loss_transducer = self.criterion_transducer( + log_probs, + target, + t_len, + u_len, + reduction="mean", + blank=self.blank_id, + fastemit_lambda=self.fastemit_lambda, + gather=True, + ) + + if not self.training and (self.report_cer or self.report_wer): + if self.error_calculator is None: + from funasr.metrics import ErrorCalculatorTransducer as ErrorCalculator + + self.error_calculator = ErrorCalculator( + self.decoder, + self.joint_network, + self.token_list, + self.sym_space, + self.sym_blank, + report_cer=self.report_cer, + report_wer=self.report_wer, + ) + + cer_transducer, wer_transducer = self.error_calculator(encoder_out, target, t_len) + + return loss_transducer, cer_transducer, wer_transducer + + return loss_transducer, None, None + + def _calc_ctc_loss( + self, + encoder_out: torch.Tensor, + target: torch.Tensor, + t_len: torch.Tensor, + u_len: torch.Tensor, + ) -> torch.Tensor: + """Compute CTC loss. + + Args: + encoder_out: Encoder output sequences. (B, T, D_enc) + target: Target label ID sequences. (B, L) + t_len: Encoder output sequences lengths. (B,) + u_len: Target label ID sequences lengths. (B,) + + Return: + loss_ctc: CTC loss value. + + """ + ctc_in = self.ctc_lin(torch.nn.functional.dropout(encoder_out, p=self.ctc_dropout_rate)) + ctc_in = torch.log_softmax(ctc_in.transpose(0, 1), dim=-1) + + target_mask = target != 0 + ctc_target = target[target_mask].cpu() + + with torch.backends.cudnn.flags(deterministic=True): + loss_ctc = torch.nn.functional.ctc_loss( + ctc_in, + ctc_target, + t_len, + u_len, + zero_infinity=True, + reduction="sum", + ) + loss_ctc /= target.size(0) + + return loss_ctc + + def _calc_lm_loss( + self, + decoder_out: torch.Tensor, + target: torch.Tensor, + ) -> torch.Tensor: + """Compute LM loss. + + Args: + decoder_out: Decoder output sequences. (B, U, D_dec) + target: Target label ID sequences. (B, L) + + Return: + loss_lm: LM loss value. + + """ + lm_loss_in = self.lm_lin(decoder_out[:, :-1, :]).view(-1, self.vocab_size) + lm_target = target.view(-1).type(torch.int64) + + with torch.no_grad(): + true_dist = lm_loss_in.clone() + true_dist.fill_(self.lm_loss_smoothing / (self.vocab_size - 1)) + + # Ignore blank ID (0) + ignore = lm_target == 0 + lm_target = lm_target.masked_fill(ignore, 0) + + true_dist.scatter_(1, lm_target.unsqueeze(1), (1 - self.lm_loss_smoothing)) + + loss_lm = torch.nn.functional.kl_div( + torch.log_softmax(lm_loss_in, dim=1), + true_dist, + reduction="none", + ) + loss_lm = loss_lm.masked_fill(ignore.unsqueeze(1), 0).sum() / decoder_out.size(0) + + return loss_lm + + def init_beam_search( + self, + **kwargs, + ): + + # 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( + length_bonus=LengthBonus(len(token_list)), + ) + + # 3. Build ngram model + # ngram is not supported now + ngram = None + scorers["ngram"] = ngram + + beam_search = BeamSearchTransducer( + self.decoder, + self.joint_network, + kwargs.get("beam_size", 2), + nbest=1, + ) + # beam_search.to(device=kwargs.get("device", "cpu"), dtype=getattr(torch, kwargs.get("dtype", "float32"))).eval() + # for scorer in scorers.values(): + # if isinstance(scorer, torch.nn.Module): + # scorer.to(device=kwargs.get("device", "cpu"), dtype=getattr(torch, kwargs.get("dtype", "float32"))).eval() + self.beam_search = beam_search + + def inference( + self, + data_in: list, + data_lengths: list = None, + key: list = None, + tokenizer=None, + **kwargs, + ): + + if kwargs.get("batch_size", 1) > 1: + raise NotImplementedError("batch decoding is not implemented") + + # init beamsearch + is_use_ctc = kwargs.get("decoding_ctc_weight", 0.0) > 0.00001 and self.ctc != None + is_use_lm = ( + kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None + ) + # if self.beam_search is None and (is_use_lm or is_use_ctc): + logging.info("enable beam_search") + self.init_beam_search(**kwargs) + self.nbest = kwargs.get("nbest", 1) + + meta_data = {} + # extract fbank feats + time1 = time.perf_counter() + audio_sample_list = load_audio_text_image_video( + data_in, fs=self.frontend.fs, audio_fs=kwargs.get("fs", 16000) + ) + time2 = time.perf_counter() + meta_data["load_data"] = f"{time2 - time1:0.3f}" + speech, speech_lengths = extract_fbank( + audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=self.frontend + ) + time3 = time.perf_counter() + meta_data["extract_feat"] = f"{time3 - time2:0.3f}" + meta_data["batch_data_time"] = ( + speech_lengths.sum().item() * self.frontend.frame_shift * self.frontend.lfr_n / 1000 + ) + + speech = speech.to(device=kwargs["device"]) + speech_lengths = speech_lengths.to(device=kwargs["device"]) + + # Encoder + encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) + if isinstance(encoder_out, tuple): + encoder_out = encoder_out[0] + + # c. Passed the encoder result and the beam search + nbest_hyps = self.beam_search(encoder_out[0], is_final=True) + nbest_hyps = nbest_hyps[: self.nbest] + + results = [] + b, n, d = encoder_out.size() + for i in range(b): + + for nbest_idx, hyp in enumerate(nbest_hyps): + ibest_writer = None + if kwargs.get("output_dir") is not None: + if not hasattr(self, "writer"): + self.writer = DatadirWriter(kwargs.get("output_dir")) + ibest_writer = self.writer[f"{nbest_idx + 1}best_recog"] + # remove sos/eos and get results + last_pos = -1 + if isinstance(hyp.yseq, list): + token_int = hyp.yseq # [1:last_pos] + else: + token_int = hyp.yseq # [1:last_pos].tolist() + + # remove blank symbol id, which is assumed to be 0 + token_int = list( + filter( + lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int + ) + ) + + # Change integer-ids to tokens + token = tokenizer.ids2tokens(token_int) + text = tokenizer.tokens2text(token) + + text_postprocessed, _ = postprocess_utils.sentence_postprocess(token) + result_i = { + "key": key[i], + "token": token, + "text": text, + "text_postprocessed": text_postprocessed, + } + results.append(result_i) + + if ibest_writer is not None: + ibest_writer["token"][key[i]] = " ".join(token) + ibest_writer["text"][key[i]] = text + ibest_writer["text_postprocessed"][key[i]] = text_postprocessed + + return results, meta_data diff --git a/funasr/models/transducer/rnn_decoder.py b/funasr/models/transducer/rnn_decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..4a0ffa3c36ccac79a6c8a22b4763bf774a03a5cb --- /dev/null +++ b/funasr/models/transducer/rnn_decoder.py @@ -0,0 +1,331 @@ +#!/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 random +import numpy as np +import torch.nn as nn +import torch.nn.functional as F + +from funasr.register import tables +from funasr.models.transformer.utils.nets_utils import make_pad_mask +from funasr.models.transformer.utils.nets_utils import to_device +from funasr.models.language_model.rnn.attentions import initial_att + + +def build_attention_list( + eprojs: int, + dunits: int, + atype: str = "location", + num_att: int = 1, + num_encs: int = 1, + aheads: int = 4, + adim: int = 320, + awin: int = 5, + aconv_chans: int = 10, + aconv_filts: int = 100, + han_mode: bool = False, + han_type=None, + han_heads: int = 4, + han_dim: int = 320, + han_conv_chans: int = -1, + han_conv_filts: int = 100, + han_win: int = 5, +): + + att_list = torch.nn.ModuleList() + if num_encs == 1: + for i in range(num_att): + att = initial_att( + atype, + eprojs, + dunits, + aheads, + adim, + awin, + aconv_chans, + aconv_filts, + ) + att_list.append(att) + elif num_encs > 1: # no multi-speaker mode + if han_mode: + att = initial_att( + han_type, + eprojs, + dunits, + han_heads, + han_dim, + han_win, + han_conv_chans, + han_conv_filts, + han_mode=True, + ) + return att + else: + att_list = torch.nn.ModuleList() + for idx in range(num_encs): + att = initial_att( + atype[idx], + eprojs, + dunits, + aheads[idx], + adim[idx], + awin[idx], + aconv_chans[idx], + aconv_filts[idx], + ) + att_list.append(att) + else: + raise ValueError("Number of encoders needs to be more than one. {}".format(num_encs)) + return att_list + + +@tables.register("decoder_classes", "rnn_decoder") +class RNNDecoder(nn.Module): + def __init__( + self, + vocab_size: int, + encoder_output_size: int, + rnn_type: str = "lstm", + num_layers: int = 1, + hidden_size: int = 320, + sampling_probability: float = 0.0, + dropout: float = 0.0, + context_residual: bool = False, + replace_sos: bool = False, + num_encs: int = 1, + att_conf: dict = None, + ): + # FIXME(kamo): The parts of num_spk should be refactored more more more + if rnn_type not in {"lstm", "gru"}: + raise ValueError(f"Not supported: rnn_type={rnn_type}") + + super().__init__() + eprojs = encoder_output_size + self.dtype = rnn_type + self.dunits = hidden_size + self.dlayers = num_layers + self.context_residual = context_residual + self.sos = vocab_size - 1 + self.eos = vocab_size - 1 + self.odim = vocab_size + self.sampling_probability = sampling_probability + self.dropout = dropout + self.num_encs = num_encs + + # for multilingual translation + self.replace_sos = replace_sos + + self.embed = torch.nn.Embedding(vocab_size, hidden_size) + self.dropout_emb = torch.nn.Dropout(p=dropout) + + self.decoder = torch.nn.ModuleList() + self.dropout_dec = torch.nn.ModuleList() + self.decoder += [ + ( + torch.nn.LSTMCell(hidden_size + eprojs, hidden_size) + if self.dtype == "lstm" + else torch.nn.GRUCell(hidden_size + eprojs, hidden_size) + ) + ] + self.dropout_dec += [torch.nn.Dropout(p=dropout)] + for _ in range(1, self.dlayers): + self.decoder += [ + ( + torch.nn.LSTMCell(hidden_size, hidden_size) + if self.dtype == "lstm" + else torch.nn.GRUCell(hidden_size, hidden_size) + ) + ] + self.dropout_dec += [torch.nn.Dropout(p=dropout)] + # NOTE: dropout is applied only for the vertical connections + # see https://arxiv.org/pdf/1409.2329.pdf + + if context_residual: + self.output = torch.nn.Linear(hidden_size + eprojs, vocab_size) + else: + self.output = torch.nn.Linear(hidden_size, vocab_size) + + self.att_list = build_attention_list(eprojs=eprojs, dunits=hidden_size, **att_conf) + + def zero_state(self, hs_pad): + return hs_pad.new_zeros(hs_pad.size(0), self.dunits) + + def rnn_forward(self, ey, z_list, c_list, z_prev, c_prev): + if self.dtype == "lstm": + z_list[0], c_list[0] = self.decoder[0](ey, (z_prev[0], c_prev[0])) + for i in range(1, self.dlayers): + z_list[i], c_list[i] = self.decoder[i]( + self.dropout_dec[i - 1](z_list[i - 1]), + (z_prev[i], c_prev[i]), + ) + else: + z_list[0] = self.decoder[0](ey, z_prev[0]) + for i in range(1, self.dlayers): + z_list[i] = self.decoder[i](self.dropout_dec[i - 1](z_list[i - 1]), z_prev[i]) + return z_list, c_list + + def forward(self, hs_pad, hlens, ys_in_pad, ys_in_lens, strm_idx=0): + # to support mutiple encoder asr mode, in single encoder mode, + # convert torch.Tensor to List of torch.Tensor + if self.num_encs == 1: + hs_pad = [hs_pad] + hlens = [hlens] + + # attention index for the attention module + # in SPA (speaker parallel attention), + # att_idx is used to select attention module. In other cases, it is 0. + att_idx = min(strm_idx, len(self.att_list) - 1) + + # hlens should be list of list of integer + hlens = [list(map(int, hlens[idx])) for idx in range(self.num_encs)] + + # get dim, length info + olength = ys_in_pad.size(1) + + # initialization + c_list = [self.zero_state(hs_pad[0])] + z_list = [self.zero_state(hs_pad[0])] + for _ in range(1, self.dlayers): + c_list.append(self.zero_state(hs_pad[0])) + z_list.append(self.zero_state(hs_pad[0])) + z_all = [] + if self.num_encs == 1: + att_w = None + self.att_list[att_idx].reset() # reset pre-computation of h + else: + att_w_list = [None] * (self.num_encs + 1) # atts + han + att_c_list = [None] * self.num_encs # atts + for idx in range(self.num_encs + 1): + # reset pre-computation of h in atts and han + self.att_list[idx].reset() + + # pre-computation of embedding + eys = self.dropout_emb(self.embed(ys_in_pad)) # utt x olen x zdim + + # loop for an output sequence + for i in range(olength): + if self.num_encs == 1: + att_c, att_w = self.att_list[att_idx]( + hs_pad[0], hlens[0], self.dropout_dec[0](z_list[0]), att_w + ) + else: + for idx in range(self.num_encs): + att_c_list[idx], att_w_list[idx] = self.att_list[idx]( + hs_pad[idx], + hlens[idx], + self.dropout_dec[0](z_list[0]), + att_w_list[idx], + ) + hs_pad_han = torch.stack(att_c_list, dim=1) + hlens_han = [self.num_encs] * len(ys_in_pad) + att_c, att_w_list[self.num_encs] = self.att_list[self.num_encs]( + hs_pad_han, + hlens_han, + self.dropout_dec[0](z_list[0]), + att_w_list[self.num_encs], + ) + if i > 0 and random.random() < self.sampling_probability: + z_out = self.output(z_all[-1]) + z_out = np.argmax(z_out.detach().cpu(), axis=1) + z_out = self.dropout_emb(self.embed(to_device(self, z_out))) + ey = torch.cat((z_out, att_c), dim=1) # utt x (zdim + hdim) + else: + # utt x (zdim + hdim) + ey = torch.cat((eys[:, i, :], att_c), dim=1) + z_list, c_list = self.rnn_forward(ey, z_list, c_list, z_list, c_list) + if self.context_residual: + z_all.append( + torch.cat((self.dropout_dec[-1](z_list[-1]), att_c), dim=-1) + ) # utt x (zdim + hdim) + else: + z_all.append(self.dropout_dec[-1](z_list[-1])) # utt x (zdim) + + z_all = torch.stack(z_all, dim=1) + z_all = self.output(z_all) + z_all.masked_fill_( + make_pad_mask(ys_in_lens, z_all, 1), + 0, + ) + return z_all, ys_in_lens + + def init_state(self, x): + # to support mutiple encoder asr mode, in single encoder mode, + # convert torch.Tensor to List of torch.Tensor + if self.num_encs == 1: + x = [x] + + c_list = [self.zero_state(x[0].unsqueeze(0))] + z_list = [self.zero_state(x[0].unsqueeze(0))] + for _ in range(1, self.dlayers): + c_list.append(self.zero_state(x[0].unsqueeze(0))) + z_list.append(self.zero_state(x[0].unsqueeze(0))) + # TODO(karita): support strm_index for `asr_mix` + strm_index = 0 + att_idx = min(strm_index, len(self.att_list) - 1) + if self.num_encs == 1: + a = None + self.att_list[att_idx].reset() # reset pre-computation of h + else: + a = [None] * (self.num_encs + 1) # atts + han + for idx in range(self.num_encs + 1): + # reset pre-computation of h in atts and han + self.att_list[idx].reset() + return dict( + c_prev=c_list[:], + z_prev=z_list[:], + a_prev=a, + workspace=(att_idx, z_list, c_list), + ) + + def score(self, yseq, state, x): + # to support mutiple encoder asr mode, in single encoder mode, + # convert torch.Tensor to List of torch.Tensor + if self.num_encs == 1: + x = [x] + + att_idx, z_list, c_list = state["workspace"] + vy = yseq[-1].unsqueeze(0) + ey = self.dropout_emb(self.embed(vy)) # utt list (1) x zdim + if self.num_encs == 1: + att_c, att_w = self.att_list[att_idx]( + x[0].unsqueeze(0), + [x[0].size(0)], + self.dropout_dec[0](state["z_prev"][0]), + state["a_prev"], + ) + else: + att_w = [None] * (self.num_encs + 1) # atts + han + att_c_list = [None] * self.num_encs # atts + for idx in range(self.num_encs): + att_c_list[idx], att_w[idx] = self.att_list[idx]( + x[idx].unsqueeze(0), + [x[idx].size(0)], + self.dropout_dec[0](state["z_prev"][0]), + state["a_prev"][idx], + ) + h_han = torch.stack(att_c_list, dim=1) + att_c, att_w[self.num_encs] = self.att_list[self.num_encs]( + h_han, + [self.num_encs], + self.dropout_dec[0](state["z_prev"][0]), + state["a_prev"][self.num_encs], + ) + ey = torch.cat((ey, att_c), dim=1) # utt(1) x (zdim + hdim) + z_list, c_list = self.rnn_forward(ey, z_list, c_list, state["z_prev"], state["c_prev"]) + if self.context_residual: + logits = self.output(torch.cat((self.dropout_dec[-1](z_list[-1]), att_c), dim=-1)) + else: + logits = self.output(self.dropout_dec[-1](z_list[-1])) + logp = F.log_softmax(logits, dim=1).squeeze(0) + return ( + logp, + dict( + c_prev=c_list[:], + z_prev=z_list[:], + a_prev=att_w, + workspace=(att_idx, z_list, c_list), + ), + ) diff --git a/funasr/models/transducer/rnnt_decoder.py b/funasr/models/transducer/rnnt_decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..bc1787ac8976aa9ece4a46e4a3ab701e851ce89b --- /dev/null +++ b/funasr/models/transducer/rnnt_decoder.py @@ -0,0 +1,264 @@ +#!/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 typing import List, Optional, Tuple + +from funasr.register import tables +from funasr.models.specaug.specaug import SpecAug +from funasr.models.transducer.beam_search_transducer import Hypothesis + + +@tables.register("decoder_classes", "rnnt_decoder") +class RNNTDecoder(torch.nn.Module): + """RNN decoder module. + + Args: + vocab_size: Vocabulary size. + embed_size: Embedding size. + hidden_size: Hidden size.. + rnn_type: Decoder layers type. + num_layers: Number of decoder layers. + dropout_rate: Dropout rate for decoder layers. + embed_dropout_rate: Dropout rate for embedding layer. + embed_pad: Embedding padding symbol ID. + + """ + + def __init__( + self, + vocab_size: int, + embed_size: int = 256, + hidden_size: int = 256, + rnn_type: str = "lstm", + num_layers: int = 1, + dropout_rate: float = 0.0, + embed_dropout_rate: float = 0.0, + embed_pad: int = 0, + use_embed_mask: bool = False, + ) -> None: + """Construct a RNNDecoder object.""" + super().__init__() + + if rnn_type not in ("lstm", "gru"): + raise ValueError(f"Not supported: rnn_type={rnn_type}") + + self.embed = torch.nn.Embedding(vocab_size, embed_size, padding_idx=embed_pad) + self.dropout_embed = torch.nn.Dropout(p=embed_dropout_rate) + + rnn_class = torch.nn.LSTM if rnn_type == "lstm" else torch.nn.GRU + + self.rnn = torch.nn.ModuleList([rnn_class(embed_size, hidden_size, 1, batch_first=True)]) + + for _ in range(1, num_layers): + self.rnn += [rnn_class(hidden_size, hidden_size, 1, batch_first=True)] + + self.dropout_rnn = torch.nn.ModuleList( + [torch.nn.Dropout(p=dropout_rate) for _ in range(num_layers)] + ) + + self.dlayers = num_layers + self.dtype = rnn_type + + self.output_size = hidden_size + self.vocab_size = vocab_size + + self.device = next(self.parameters()).device + self.score_cache = {} + + self.use_embed_mask = use_embed_mask + if self.use_embed_mask: + self._embed_mask = SpecAug( + time_mask_width_range=3, + num_time_mask=4, + apply_freq_mask=False, + apply_time_warp=False, + ) + + def forward( + self, + labels: torch.Tensor, + label_lens: torch.Tensor, + states: Optional[Tuple[torch.Tensor, Optional[torch.Tensor]]] = None, + ) -> torch.Tensor: + """Encode source label sequences. + + Args: + labels: Label ID sequences. (B, L) + states: Decoder hidden states. + ((N, B, D_dec), (N, B, D_dec) or None) or None + + Returns: + dec_out: Decoder output sequences. (B, U, D_dec) + + """ + if states is None: + states = self.init_state(labels.size(0)) + + dec_embed = self.dropout_embed(self.embed(labels)) + if self.use_embed_mask and self.training: + dec_embed = self._embed_mask(dec_embed, label_lens)[0] + dec_out, states = self.rnn_forward(dec_embed, states) + return dec_out + + def rnn_forward( + self, + x: torch.Tensor, + state: Tuple[torch.Tensor, Optional[torch.Tensor]], + ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, Optional[torch.Tensor]]]: + """Encode source label sequences. + + Args: + x: RNN input sequences. (B, D_emb) + state: Decoder hidden states. ((N, B, D_dec), (N, B, D_dec) or None) + + Returns: + x: RNN output sequences. (B, D_dec) + (h_next, c_next): Decoder hidden states. + (N, B, D_dec), (N, B, D_dec) or None) + + """ + h_prev, c_prev = state + h_next, c_next = self.init_state(x.size(0)) + + for layer in range(self.dlayers): + if self.dtype == "lstm": + x, (h_next[layer : layer + 1], c_next[layer : layer + 1]) = self.rnn[layer]( + x, hx=(h_prev[layer : layer + 1], c_prev[layer : layer + 1]) + ) + else: + x, h_next[layer : layer + 1] = self.rnn[layer](x, hx=h_prev[layer : layer + 1]) + + x = self.dropout_rnn[layer](x) + + return x, (h_next, c_next) + + def score( + self, + label: torch.Tensor, + label_sequence: List[int], + dec_state: Tuple[torch.Tensor, Optional[torch.Tensor]], + ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, Optional[torch.Tensor]]]: + """One-step forward hypothesis. + + Args: + label: Previous label. (1, 1) + label_sequence: Current label sequence. + dec_state: Previous decoder hidden states. + ((N, 1, D_dec), (N, 1, D_dec) or None) + + Returns: + dec_out: Decoder output sequence. (1, D_dec) + dec_state: Decoder hidden states. + ((N, 1, D_dec), (N, 1, D_dec) or None) + + """ + str_labels = "_".join(map(str, label_sequence)) + + if str_labels in self.score_cache: + dec_out, dec_state = self.score_cache[str_labels] + else: + dec_embed = self.embed(label) + dec_out, dec_state = self.rnn_forward(dec_embed, dec_state) + + self.score_cache[str_labels] = (dec_out, dec_state) + + return dec_out[0], dec_state + + def batch_score( + self, + hyps: List[Hypothesis], + ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, Optional[torch.Tensor]]]: + """One-step forward hypotheses. + + Args: + hyps: Hypotheses. + + Returns: + dec_out: Decoder output sequences. (B, D_dec) + states: Decoder hidden states. ((N, B, D_dec), (N, B, D_dec) or None) + + """ + labels = torch.LongTensor([[h.yseq[-1]] for h in hyps], device=self.device) + dec_embed = self.embed(labels) + + states = self.create_batch_states([h.dec_state for h in hyps]) + dec_out, states = self.rnn_forward(dec_embed, states) + + return dec_out.squeeze(1), states + + def set_device(self, device: torch.device) -> None: + """Set GPU device to use. + + Args: + device: Device ID. + + """ + self.device = device + + def init_state(self, batch_size: int) -> Tuple[torch.Tensor, Optional[torch.tensor]]: + """Initialize decoder states. + + Args: + batch_size: Batch size. + + Returns: + : Initial decoder hidden states. ((N, B, D_dec), (N, B, D_dec) or None) + + """ + h_n = torch.zeros( + self.dlayers, + batch_size, + self.output_size, + device=self.device, + ) + + if self.dtype == "lstm": + c_n = torch.zeros( + self.dlayers, + batch_size, + self.output_size, + device=self.device, + ) + + return (h_n, c_n) + + return (h_n, None) + + def select_state( + self, states: Tuple[torch.Tensor, Optional[torch.Tensor]], idx: int + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Get specified ID state from decoder hidden states. + + Args: + states: Decoder hidden states. ((N, B, D_dec), (N, B, D_dec) or None) + idx: State ID to extract. + + Returns: + : Decoder hidden state for given ID. ((N, 1, D_dec), (N, 1, D_dec) or None) + + """ + return ( + states[0][:, idx : idx + 1, :], + states[1][:, idx : idx + 1, :] if self.dtype == "lstm" else None, + ) + + def create_batch_states( + self, + new_states: List[Tuple[torch.Tensor, Optional[torch.Tensor]]], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Create decoder hidden states. + + Args: + new_states: Decoder hidden states. [N x ((1, D_dec), (1, D_dec) or None)] + + Returns: + states: Decoder hidden states. ((N, B, D_dec), (N, B, D_dec) or None) + + """ + return ( + torch.cat([s[0] for s in new_states], dim=1), + torch.cat([s[1] for s in new_states], dim=1) if self.dtype == "lstm" else None, + ) diff --git a/funasr/models/transformer/__init__.py b/funasr/models/transformer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/models/transformer/attention.py b/funasr/models/transformer/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..2844333a12e39abae698a3dadb3323bc42647d6e --- /dev/null +++ b/funasr/models/transformer/attention.py @@ -0,0 +1,613 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +# Copyright 2019 Shigeki Karita +# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +"""Multi-Head Attention layer definition.""" + +import math + +import numpy +import torch +from torch import nn +from typing import Optional, Tuple + +import torch.nn.functional as F +from funasr.models.transformer.utils.nets_utils import make_pad_mask +import funasr.models.lora.layers as lora + + +class MultiHeadedAttention(nn.Module): + """Multi-Head Attention layer. + + Args: + n_head (int): The number of heads. + n_feat (int): The number of features. + dropout_rate (float): Dropout rate. + + """ + + def __init__(self, n_head, n_feat, dropout_rate): + """Construct an MultiHeadedAttention object.""" + super(MultiHeadedAttention, self).__init__() + assert n_feat % n_head == 0 + # We assume d_v always equals d_k + self.d_k = n_feat // n_head + self.h = n_head + self.linear_q = nn.Linear(n_feat, n_feat) + self.linear_k = nn.Linear(n_feat, n_feat) + self.linear_v = nn.Linear(n_feat, n_feat) + self.linear_out = nn.Linear(n_feat, n_feat) + self.attn = None + self.dropout = nn.Dropout(p=dropout_rate) + + def forward_qkv(self, query, key, value): + """Transform query, key and value. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + + Returns: + torch.Tensor: Transformed query tensor (#batch, n_head, time1, d_k). + torch.Tensor: Transformed key tensor (#batch, n_head, time2, d_k). + torch.Tensor: Transformed value tensor (#batch, n_head, time2, d_k). + + """ + n_batch = query.size(0) + q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k) + k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k) + v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k) + q = q.transpose(1, 2) # (batch, head, time1, d_k) + k = k.transpose(1, 2) # (batch, head, time2, d_k) + v = v.transpose(1, 2) # (batch, head, time2, d_k) + + return q, k, v + + def forward_attention(self, value, scores, mask): + """Compute attention context vector. + + Args: + value (torch.Tensor): Transformed value (#batch, n_head, time2, d_k). + scores (torch.Tensor): Attention score (#batch, n_head, time1, time2). + mask (torch.Tensor): Mask (#batch, 1, time2) or (#batch, time1, time2). + + Returns: + torch.Tensor: Transformed value (#batch, time1, d_model) + weighted by the attention score (#batch, time1, time2). + + """ + n_batch = value.size(0) + if mask is not None: + mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) + + min_value = -float( + "inf" + ) # min_value = float(np.finfo(torch.tensor(0, dtype=qk.dtype).numpy().dtype).min) + scores = scores.masked_fill(mask, min_value) + attn = torch.softmax(scores, dim=-1).masked_fill( + mask, 0.0 + ) # (batch, head, time1, time2) + else: + attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2) + + p_attn = self.dropout(attn) + x = torch.matmul(p_attn, value) # (batch, head, time1, d_k) + x = ( + x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) + ) # (batch, time1, d_model) + + return self.linear_out(x) # (batch, time1, d_model) + + def forward(self, query, key, value, mask): + """Compute scaled dot product attention. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + mask (torch.Tensor): Mask tensor (#batch, 1, time2) or + (#batch, time1, time2). + + Returns: + torch.Tensor: Output tensor (#batch, time1, d_model). + + """ + q, k, v = self.forward_qkv(query, key, value) + scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) + return self.forward_attention(v, scores, mask) + + +class MultiHeadedAttentionExport(nn.Module): + def __init__(self, model): + super().__init__() + self.d_k = model.d_k + self.h = model.h + self.linear_q = model.linear_q + self.linear_k = model.linear_k + self.linear_v = model.linear_v + self.linear_out = model.linear_out + self.attn = None + self.all_head_size = self.h * self.d_k + + def forward(self, query, key, value, mask): + q, k, v = self.forward_qkv(query, key, value) + scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) + return self.forward_attention(v, scores, mask) + + def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor: + new_x_shape = x.size()[:-1] + (self.h, self.d_k) + x = x.view(new_x_shape) + return x.permute(0, 2, 1, 3) + + def forward_qkv(self, query, key, value): + q = self.linear_q(query) + k = self.linear_k(key) + v = self.linear_v(value) + q = self.transpose_for_scores(q) + k = self.transpose_for_scores(k) + v = self.transpose_for_scores(v) + return q, k, v + + def forward_attention(self, value, scores, mask): + scores = scores + mask + + attn = torch.softmax(scores, dim=-1) + context_layer = torch.matmul(attn, value) # (batch, head, time1, d_k) + + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.view(new_context_layer_shape) + return self.linear_out(context_layer) # (batch, time1, d_model) + + +class RelPosMultiHeadedAttentionExport(MultiHeadedAttentionExport): + def __init__(self, model): + super().__init__(model) + self.linear_pos = model.linear_pos + self.pos_bias_u = model.pos_bias_u + self.pos_bias_v = model.pos_bias_v + + def forward(self, query, key, value, pos_emb, mask): + q, k, v = self.forward_qkv(query, key, value) + q = q.transpose(1, 2) # (batch, time1, head, d_k) + + p = self.transpose_for_scores(self.linear_pos(pos_emb)) # (batch, head, time1, d_k) + + # (batch, head, time1, d_k) + q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2) + # (batch, head, time1, d_k) + q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2) + + # compute attention score + # first compute matrix a and matrix c + # as described in https://arxiv.org/abs/1901.02860 Section 3.3 + # (batch, head, time1, time2) + matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1)) + + # compute matrix b and matrix d + # (batch, head, time1, time1) + matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) + matrix_bd = self.rel_shift(matrix_bd) + + scores = (matrix_ac + matrix_bd) / math.sqrt(self.d_k) # (batch, head, time1, time2) + + return self.forward_attention(v, scores, mask) + + def rel_shift(self, x): + zero_pad = torch.zeros((*x.size()[:3], 1), device=x.device, dtype=x.dtype) + x_padded = torch.cat([zero_pad, x], dim=-1) + + x_padded = x_padded.view(*x.size()[:2], x.size(3) + 1, x.size(2)) + x = x_padded[:, :, 1:].view_as(x)[ + :, :, :, : x.size(-1) // 2 + 1 + ] # only keep the positions from 0 to time2 + return x + + def forward_attention(self, value, scores, mask): + scores = scores + mask + + attn = torch.softmax(scores, dim=-1) + context_layer = torch.matmul(attn, value) # (batch, head, time1, d_k) + + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.view(new_context_layer_shape) + return self.linear_out(context_layer) # (batch, time1, d_model) + + +class LegacyRelPositionMultiHeadedAttention(MultiHeadedAttention): + """Multi-Head Attention layer with relative position encoding (old version). + + Details can be found in https://github.com/espnet/espnet/pull/2816. + + Paper: https://arxiv.org/abs/1901.02860 + + Args: + n_head (int): The number of heads. + n_feat (int): The number of features. + dropout_rate (float): Dropout rate. + zero_triu (bool): Whether to zero the upper triangular part of attention matrix. + + """ + + def __init__(self, n_head, n_feat, dropout_rate, zero_triu=False): + """Construct an RelPositionMultiHeadedAttention object.""" + super().__init__(n_head, n_feat, dropout_rate) + self.zero_triu = zero_triu + # linear transformation for positional encoding + self.linear_pos = nn.Linear(n_feat, n_feat, bias=False) + # these two learnable bias are used in matrix c and matrix d + # as described in https://arxiv.org/abs/1901.02860 Section 3.3 + self.pos_bias_u = nn.Parameter(torch.Tensor(self.h, self.d_k)) + self.pos_bias_v = nn.Parameter(torch.Tensor(self.h, self.d_k)) + torch.nn.init.xavier_uniform_(self.pos_bias_u) + torch.nn.init.xavier_uniform_(self.pos_bias_v) + + def rel_shift(self, x): + """Compute relative positional encoding. + + Args: + x (torch.Tensor): Input tensor (batch, head, time1, time2). + + Returns: + torch.Tensor: Output tensor. + + """ + zero_pad = torch.zeros((*x.size()[:3], 1), device=x.device, dtype=x.dtype) + x_padded = torch.cat([zero_pad, x], dim=-1) + + x_padded = x_padded.view(*x.size()[:2], x.size(3) + 1, x.size(2)) + x = x_padded[:, :, 1:].view_as(x) + + if self.zero_triu: + ones = torch.ones((x.size(2), x.size(3))) + x = x * torch.tril(ones, x.size(3) - x.size(2))[None, None, :, :] + + return x + + def forward(self, query, key, value, pos_emb, mask): + """Compute 'Scaled Dot Product Attention' with rel. positional encoding. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + pos_emb (torch.Tensor): Positional embedding tensor (#batch, time1, size). + mask (torch.Tensor): Mask tensor (#batch, 1, time2) or + (#batch, time1, time2). + + Returns: + torch.Tensor: Output tensor (#batch, time1, d_model). + + """ + q, k, v = self.forward_qkv(query, key, value) + q = q.transpose(1, 2) # (batch, time1, head, d_k) + + n_batch_pos = pos_emb.size(0) + p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k) + p = p.transpose(1, 2) # (batch, head, time1, d_k) + + # (batch, head, time1, d_k) + q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2) + # (batch, head, time1, d_k) + q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2) + + # compute attention score + # first compute matrix a and matrix c + # as described in https://arxiv.org/abs/1901.02860 Section 3.3 + # (batch, head, time1, time2) + matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1)) + + # compute matrix b and matrix d + # (batch, head, time1, time1) + matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) + matrix_bd = self.rel_shift(matrix_bd) + + scores = (matrix_ac + matrix_bd) / math.sqrt(self.d_k) # (batch, head, time1, time2) + + return self.forward_attention(v, scores, mask) + + +class RelPositionMultiHeadedAttention(MultiHeadedAttention): + """Multi-Head Attention layer with relative position encoding (new implementation). + + Details can be found in https://github.com/espnet/espnet/pull/2816. + + Paper: https://arxiv.org/abs/1901.02860 + + Args: + n_head (int): The number of heads. + n_feat (int): The number of features. + dropout_rate (float): Dropout rate. + zero_triu (bool): Whether to zero the upper triangular part of attention matrix. + + """ + + def __init__(self, n_head, n_feat, dropout_rate, zero_triu=False): + """Construct an RelPositionMultiHeadedAttention object.""" + super().__init__(n_head, n_feat, dropout_rate) + self.zero_triu = zero_triu + # linear transformation for positional encoding + self.linear_pos = nn.Linear(n_feat, n_feat, bias=False) + # these two learnable bias are used in matrix c and matrix d + # as described in https://arxiv.org/abs/1901.02860 Section 3.3 + self.pos_bias_u = nn.Parameter(torch.Tensor(self.h, self.d_k)) + self.pos_bias_v = nn.Parameter(torch.Tensor(self.h, self.d_k)) + torch.nn.init.xavier_uniform_(self.pos_bias_u) + torch.nn.init.xavier_uniform_(self.pos_bias_v) + + def rel_shift(self, x): + """Compute relative positional encoding. + + Args: + x (torch.Tensor): Input tensor (batch, head, time1, 2*time1-1). + time1 means the length of query vector. + + Returns: + torch.Tensor: Output tensor. + + """ + zero_pad = torch.zeros((*x.size()[:3], 1), device=x.device, dtype=x.dtype) + x_padded = torch.cat([zero_pad, x], dim=-1) + + x_padded = x_padded.view(*x.size()[:2], x.size(3) + 1, x.size(2)) + x = x_padded[:, :, 1:].view_as(x)[ + :, :, :, : x.size(-1) // 2 + 1 + ] # only keep the positions from 0 to time2 + + if self.zero_triu: + ones = torch.ones((x.size(2), x.size(3)), device=x.device) + x = x * torch.tril(ones, x.size(3) - x.size(2))[None, None, :, :] + + return x + + def forward(self, query, key, value, pos_emb, mask): + """Compute 'Scaled Dot Product Attention' with rel. positional encoding. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + pos_emb (torch.Tensor): Positional embedding tensor + (#batch, 2*time1-1, size). + mask (torch.Tensor): Mask tensor (#batch, 1, time2) or + (#batch, time1, time2). + + Returns: + torch.Tensor: Output tensor (#batch, time1, d_model). + + """ + q, k, v = self.forward_qkv(query, key, value) + q = q.transpose(1, 2) # (batch, time1, head, d_k) + + n_batch_pos = pos_emb.size(0) + p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k) + p = p.transpose(1, 2) # (batch, head, 2*time1-1, d_k) + + # (batch, head, time1, d_k) + q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2) + # (batch, head, time1, d_k) + q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2) + + # compute attention score + # first compute matrix a and matrix c + # as described in https://arxiv.org/abs/1901.02860 Section 3.3 + # (batch, head, time1, time2) + matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1)) + + # compute matrix b and matrix d + # (batch, head, time1, 2*time1-1) + matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) + matrix_bd = self.rel_shift(matrix_bd) + + scores = (matrix_ac + matrix_bd) / math.sqrt(self.d_k) # (batch, head, time1, time2) + + return self.forward_attention(v, scores, mask) + + +class RelPositionMultiHeadedAttentionChunk(torch.nn.Module): + """RelPositionMultiHeadedAttention definition. + Args: + num_heads: Number of attention heads. + embed_size: Embedding size. + dropout_rate: Dropout rate. + """ + + def __init__( + self, + num_heads: int, + embed_size: int, + dropout_rate: float = 0.0, + simplified_attention_score: bool = False, + ) -> None: + """Construct an MultiHeadedAttention object.""" + super().__init__() + + self.d_k = embed_size // num_heads + self.num_heads = num_heads + + assert self.d_k * num_heads == embed_size, ( + "embed_size (%d) must be divisible by num_heads (%d)", + (embed_size, num_heads), + ) + + self.linear_q = torch.nn.Linear(embed_size, embed_size) + self.linear_k = torch.nn.Linear(embed_size, embed_size) + self.linear_v = torch.nn.Linear(embed_size, embed_size) + + self.linear_out = torch.nn.Linear(embed_size, embed_size) + + if simplified_attention_score: + self.linear_pos = torch.nn.Linear(embed_size, num_heads) + + self.compute_att_score = self.compute_simplified_attention_score + else: + self.linear_pos = torch.nn.Linear(embed_size, embed_size, bias=False) + + self.pos_bias_u = torch.nn.Parameter(torch.Tensor(num_heads, self.d_k)) + self.pos_bias_v = torch.nn.Parameter(torch.Tensor(num_heads, self.d_k)) + torch.nn.init.xavier_uniform_(self.pos_bias_u) + torch.nn.init.xavier_uniform_(self.pos_bias_v) + + self.compute_att_score = self.compute_attention_score + + self.dropout = torch.nn.Dropout(p=dropout_rate) + self.attn = None + + def rel_shift(self, x: torch.Tensor, left_context: int = 0) -> torch.Tensor: + """Compute relative positional encoding. + Args: + x: Input sequence. (B, H, T_1, 2 * T_1 - 1) + left_context: Number of frames in left context. + Returns: + x: Output sequence. (B, H, T_1, T_2) + """ + batch_size, n_heads, time1, n = x.shape + time2 = time1 + left_context + + batch_stride, n_heads_stride, time1_stride, n_stride = x.stride() + + return x.as_strided( + (batch_size, n_heads, time1, time2), + (batch_stride, n_heads_stride, time1_stride - n_stride, n_stride), + storage_offset=(n_stride * (time1 - 1)), + ) + + def compute_simplified_attention_score( + self, + query: torch.Tensor, + key: torch.Tensor, + pos_enc: torch.Tensor, + left_context: int = 0, + ) -> torch.Tensor: + """Simplified attention score computation. + Reference: https://github.com/k2-fsa/icefall/pull/458 + Args: + query: Transformed query tensor. (B, H, T_1, d_k) + key: Transformed key tensor. (B, H, T_2, d_k) + pos_enc: Positional embedding tensor. (B, 2 * T_1 - 1, size) + left_context: Number of frames in left context. + Returns: + : Attention score. (B, H, T_1, T_2) + """ + pos_enc = self.linear_pos(pos_enc) + + matrix_ac = torch.matmul(query, key.transpose(2, 3)) + + matrix_bd = self.rel_shift( + pos_enc.transpose(1, 2).unsqueeze(2).repeat(1, 1, query.size(2), 1), + left_context=left_context, + ) + + return (matrix_ac + matrix_bd) / math.sqrt(self.d_k) + + def compute_attention_score( + self, + query: torch.Tensor, + key: torch.Tensor, + pos_enc: torch.Tensor, + left_context: int = 0, + ) -> torch.Tensor: + """Attention score computation. + Args: + query: Transformed query tensor. (B, H, T_1, d_k) + key: Transformed key tensor. (B, H, T_2, d_k) + pos_enc: Positional embedding tensor. (B, 2 * T_1 - 1, size) + left_context: Number of frames in left context. + Returns: + : Attention score. (B, H, T_1, T_2) + """ + p = self.linear_pos(pos_enc).view(pos_enc.size(0), -1, self.num_heads, self.d_k) + + query = query.transpose(1, 2) + q_with_bias_u = (query + self.pos_bias_u).transpose(1, 2) + q_with_bias_v = (query + self.pos_bias_v).transpose(1, 2) + + matrix_ac = torch.matmul(q_with_bias_u, key.transpose(-2, -1)) + + matrix_bd = torch.matmul(q_with_bias_v, p.permute(0, 2, 3, 1)) + matrix_bd = self.rel_shift(matrix_bd, left_context=left_context) + + return (matrix_ac + matrix_bd) / math.sqrt(self.d_k) + + def forward_qkv( + self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Transform query, key and value. + Args: + query: Query tensor. (B, T_1, size) + key: Key tensor. (B, T_2, size) + v: Value tensor. (B, T_2, size) + Returns: + q: Transformed query tensor. (B, H, T_1, d_k) + k: Transformed key tensor. (B, H, T_2, d_k) + v: Transformed value tensor. (B, H, T_2, d_k) + """ + n_batch = query.size(0) + + q = self.linear_q(query).view(n_batch, -1, self.num_heads, self.d_k).transpose(1, 2) + k = self.linear_k(key).view(n_batch, -1, self.num_heads, self.d_k).transpose(1, 2) + v = self.linear_v(value).view(n_batch, -1, self.num_heads, self.d_k).transpose(1, 2) + + return q, k, v + + def forward_attention( + self, + value: torch.Tensor, + scores: torch.Tensor, + mask: torch.Tensor, + chunk_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Compute attention context vector. + Args: + value: Transformed value. (B, H, T_2, d_k) + scores: Attention score. (B, H, T_1, T_2) + mask: Source mask. (B, T_2) + chunk_mask: Chunk mask. (T_1, T_1) + Returns: + attn_output: Transformed value weighted by attention score. (B, T_1, H * d_k) + """ + batch_size = scores.size(0) + mask = mask.unsqueeze(1).unsqueeze(2) + if chunk_mask is not None: + mask = chunk_mask.unsqueeze(0).unsqueeze(1) | mask + scores = scores.masked_fill(mask, float("-inf")) + attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0) + + attn_output = self.dropout(attn) + attn_output = torch.matmul(attn_output, value) + + attn_output = self.linear_out( + attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.d_k) + ) + + return attn_output + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + pos_enc: torch.Tensor, + mask: torch.Tensor, + chunk_mask: Optional[torch.Tensor] = None, + left_context: int = 0, + ) -> torch.Tensor: + """Compute scaled dot product attention with rel. positional encoding. + Args: + query: Query tensor. (B, T_1, size) + key: Key tensor. (B, T_2, size) + value: Value tensor. (B, T_2, size) + pos_enc: Positional embedding tensor. (B, 2 * T_1 - 1, size) + mask: Source mask. (B, T_2) + chunk_mask: Chunk mask. (T_1, T_1) + left_context: Number of frames in left context. + Returns: + : Output tensor. (B, T_1, H * d_k) + """ + q, k, v = self.forward_qkv(query, key, value) + scores = self.compute_att_score(q, k, pos_enc, left_context=left_context) + return self.forward_attention(v, scores, mask, chunk_mask=chunk_mask) diff --git a/funasr/models/transformer/decoder.py b/funasr/models/transformer/decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..7d849f3498e96319d0f68536c2203b34b8cb30be --- /dev/null +++ b/funasr/models/transformer/decoder.py @@ -0,0 +1,653 @@ +# Copyright 2019 Shigeki Karita +# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +"""Decoder definition.""" +from typing import Any +from typing import List +from typing import Sequence +from typing import Tuple + +import torch +from torch import nn + + +from funasr.models.transformer.attention import MultiHeadedAttention +from funasr.models.transformer.utils.dynamic_conv import DynamicConvolution +from funasr.models.transformer.utils.dynamic_conv2d import DynamicConvolution2D +from funasr.models.transformer.embedding import PositionalEncoding +from funasr.models.transformer.layer_norm import LayerNorm +from funasr.models.transformer.utils.lightconv import LightweightConvolution +from funasr.models.transformer.utils.lightconv2d import LightweightConvolution2D +from funasr.models.transformer.utils.mask import subsequent_mask +from funasr.models.transformer.utils.nets_utils import make_pad_mask +from funasr.models.transformer.positionwise_feed_forward import ( + PositionwiseFeedForward, # noqa: H301 +) +from funasr.models.transformer.utils.repeat import repeat +from funasr.models.transformer.scorers.scorer_interface import BatchScorerInterface + +from funasr.register import tables + + +class DecoderLayer(nn.Module): + """Single decoder layer module. + + Args: + size (int): Input dimension. + self_attn (torch.nn.Module): Self-attention module instance. + `MultiHeadedAttention` instance can be used as the argument. + src_attn (torch.nn.Module): Self-attention module instance. + `MultiHeadedAttention` instance can be used as the argument. + feed_forward (torch.nn.Module): Feed-forward module instance. + `PositionwiseFeedForward`, `MultiLayeredConv1d`, or `Conv1dLinear` instance + can be used as the argument. + dropout_rate (float): Dropout rate. + normalize_before (bool): Whether to use layer_norm before the first block. + concat_after (bool): Whether to concat attention layer's input and output. + if True, additional linear will be applied. + i.e. x -> x + linear(concat(x, att(x))) + if False, no additional linear will be applied. i.e. x -> x + att(x) + + + """ + + def __init__( + self, + size, + self_attn, + src_attn, + feed_forward, + dropout_rate, + normalize_before=True, + concat_after=False, + ): + """Construct an DecoderLayer object.""" + super(DecoderLayer, self).__init__() + self.size = size + self.self_attn = self_attn + self.src_attn = src_attn + self.feed_forward = feed_forward + self.norm1 = LayerNorm(size) + self.norm2 = LayerNorm(size) + self.norm3 = LayerNorm(size) + self.dropout = nn.Dropout(dropout_rate) + self.normalize_before = normalize_before + self.concat_after = concat_after + if self.concat_after: + self.concat_linear1 = nn.Linear(size + size, size) + self.concat_linear2 = nn.Linear(size + size, size) + + def forward(self, tgt, tgt_mask, memory, memory_mask, cache=None): + """Compute decoded features. + + Args: + tgt (torch.Tensor): Input tensor (#batch, maxlen_out, size). + tgt_mask (torch.Tensor): Mask for input tensor (#batch, maxlen_out). + memory (torch.Tensor): Encoded memory, float32 (#batch, maxlen_in, size). + memory_mask (torch.Tensor): Encoded memory mask (#batch, maxlen_in). + cache (List[torch.Tensor]): List of cached tensors. + Each tensor shape should be (#batch, maxlen_out - 1, size). + + Returns: + torch.Tensor: Output tensor(#batch, maxlen_out, size). + torch.Tensor: Mask for output tensor (#batch, maxlen_out). + torch.Tensor: Encoded memory (#batch, maxlen_in, size). + torch.Tensor: Encoded memory mask (#batch, maxlen_in). + + """ + residual = tgt + if self.normalize_before: + tgt = self.norm1(tgt) + + if cache is None: + tgt_q = tgt + tgt_q_mask = tgt_mask + else: + # compute only the last frame query keeping dim: max_time_out -> 1 + assert cache.shape == ( + tgt.shape[0], + tgt.shape[1] - 1, + self.size, + ), f"{cache.shape} == {(tgt.shape[0], tgt.shape[1] - 1, self.size)}" + tgt_q = tgt[:, -1:, :] + residual = residual[:, -1:, :] + tgt_q_mask = None + if tgt_mask is not None: + tgt_q_mask = tgt_mask[:, -1:, :] + + if self.concat_after: + tgt_concat = torch.cat((tgt_q, self.self_attn(tgt_q, tgt, tgt, tgt_q_mask)), dim=-1) + x = residual + self.concat_linear1(tgt_concat) + else: + x = residual + self.dropout(self.self_attn(tgt_q, tgt, tgt, tgt_q_mask)) + if not self.normalize_before: + x = self.norm1(x) + + residual = x + if self.normalize_before: + x = self.norm2(x) + if self.concat_after: + x_concat = torch.cat((x, self.src_attn(x, memory, memory, memory_mask)), dim=-1) + x = residual + self.concat_linear2(x_concat) + else: + x = residual + self.dropout(self.src_attn(x, memory, memory, memory_mask)) + if not self.normalize_before: + x = self.norm2(x) + + residual = x + if self.normalize_before: + x = self.norm3(x) + x = residual + self.dropout(self.feed_forward(x)) + if not self.normalize_before: + x = self.norm3(x) + + if cache is not None: + x = torch.cat([cache, x], dim=1) + + return x, tgt_mask, memory, memory_mask + + +class DecoderLayerExport(nn.Module): + def __init__(self, model): + super().__init__() + self.self_attn = model.self_attn + self.src_attn = model.src_attn + self.feed_forward = model.feed_forward + self.norm1 = model.norm1 + self.norm2 = model.norm2 + self.norm3 = model.norm3 + + def forward(self, tgt, tgt_mask, memory, memory_mask, cache=None): + residual = tgt + tgt = self.norm1(tgt) + tgt_q = tgt + tgt_q_mask = tgt_mask + x = residual + self.self_attn(tgt_q, tgt, tgt, tgt_q_mask) + + residual = x + x = self.norm2(x) + + x = residual + self.src_attn(x, memory, memory, memory_mask) + + residual = x + x = self.norm3(x) + x = residual + self.feed_forward(x) + + return x, tgt_mask, memory, memory_mask + + +class BaseTransformerDecoder(nn.Module, BatchScorerInterface): + """Base class of Transfomer decoder module. + + Args: + vocab_size: output dim + encoder_output_size: dimension of attention + attention_heads: the number of heads of multi head attention + linear_units: the number of units of position-wise feed forward + num_blocks: the number of decoder blocks + dropout_rate: dropout rate + self_attention_dropout_rate: dropout rate for attention + input_layer: input layer type + use_output_layer: whether to use output layer + pos_enc_class: PositionalEncoding or ScaledPositionalEncoding + normalize_before: whether to use layer_norm before the first block + concat_after: whether to concat attention layer's input and output + if True, additional linear will be applied. + i.e. x -> x + linear(concat(x, att(x))) + if False, no additional linear will be applied. + i.e. x -> x + att(x) + """ + + def __init__( + self, + vocab_size: int, + encoder_output_size: int, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + ): + super().__init__() + attention_dim = encoder_output_size + + if input_layer == "embed": + self.embed = torch.nn.Sequential( + torch.nn.Embedding(vocab_size, attention_dim), + pos_enc_class(attention_dim, positional_dropout_rate), + ) + elif input_layer == "linear": + self.embed = torch.nn.Sequential( + torch.nn.Linear(vocab_size, attention_dim), + torch.nn.LayerNorm(attention_dim), + torch.nn.Dropout(dropout_rate), + torch.nn.ReLU(), + pos_enc_class(attention_dim, positional_dropout_rate), + ) + else: + raise ValueError(f"only 'embed' or 'linear' is supported: {input_layer}") + + self.normalize_before = normalize_before + if self.normalize_before: + self.after_norm = LayerNorm(attention_dim) + if use_output_layer: + self.output_layer = torch.nn.Linear(attention_dim, vocab_size) + else: + self.output_layer = None + + # Must set by the inheritance + self.decoders = None + + def forward( + self, + hs_pad: torch.Tensor, + hlens: torch.Tensor, + ys_in_pad: torch.Tensor, + ys_in_lens: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Forward decoder. + + Args: + hs_pad: encoded memory, float32 (batch, maxlen_in, feat) + hlens: (batch) + ys_in_pad: + input token ids, int64 (batch, maxlen_out) + if input_layer == "embed" + input tensor (batch, maxlen_out, #mels) in the other cases + ys_in_lens: (batch) + Returns: + (tuple): tuple containing: + + x: decoded token score before softmax (batch, maxlen_out, token) + if use_output_layer is True, + olens: (batch, ) + """ + tgt = ys_in_pad + # tgt_mask: (B, 1, L) + tgt_mask = (~make_pad_mask(ys_in_lens)[:, None, :]).to(tgt.device) + # m: (1, L, L) + m = subsequent_mask(tgt_mask.size(-1), device=tgt_mask.device).unsqueeze(0) + # tgt_mask: (B, L, L) + tgt_mask = tgt_mask & m + + memory = hs_pad + memory_mask = (~make_pad_mask(hlens, maxlen=memory.size(1)))[:, None, :].to(memory.device) + # Padding for Longformer + if memory_mask.shape[-1] != memory.shape[1]: + padlen = memory.shape[1] - memory_mask.shape[-1] + memory_mask = torch.nn.functional.pad(memory_mask, (0, padlen), "constant", False) + + x = self.embed(tgt) + x, tgt_mask, memory, memory_mask = self.decoders(x, tgt_mask, memory, memory_mask) + if self.normalize_before: + x = self.after_norm(x) + if self.output_layer is not None: + x = self.output_layer(x) + + olens = tgt_mask.sum(1) + return x, olens + + def forward_one_step( + self, + tgt: torch.Tensor, + tgt_mask: torch.Tensor, + memory: torch.Tensor, + cache: List[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, List[torch.Tensor]]: + """Forward one step. + + Args: + tgt: input token ids, int64 (batch, maxlen_out) + tgt_mask: input token mask, (batch, maxlen_out) + dtype=torch.uint8 in PyTorch 1.2- + dtype=torch.bool in PyTorch 1.2+ (include 1.2) + memory: encoded memory, float32 (batch, maxlen_in, feat) + cache: cached output list of (batch, max_time_out-1, size) + Returns: + y, cache: NN output value and cache per `self.decoders`. + y.shape` is (batch, maxlen_out, token) + """ + x = self.embed(tgt) + if cache is None: + cache = [None] * len(self.decoders) + new_cache = [] + for c, decoder in zip(cache, self.decoders): + x, tgt_mask, memory, memory_mask = decoder(x, tgt_mask, memory, None, cache=c) + new_cache.append(x) + + if self.normalize_before: + y = self.after_norm(x[:, -1]) + else: + y = x[:, -1] + if self.output_layer is not None: + y = torch.log_softmax(self.output_layer(y), dim=-1) + + return y, new_cache + + def score(self, ys, state, x): + """Score.""" + ys_mask = subsequent_mask(len(ys), device=x.device).unsqueeze(0) + logp, state = self.forward_one_step(ys.unsqueeze(0), ys_mask, x.unsqueeze(0), cache=state) + return logp.squeeze(0), state + + def batch_score( + self, ys: torch.Tensor, states: List[Any], xs: torch.Tensor + ) -> Tuple[torch.Tensor, List[Any]]: + """Score new token batch. + + Args: + ys (torch.Tensor): torch.int64 prefix tokens (n_batch, ylen). + states (List[Any]): Scorer states for prefix tokens. + xs (torch.Tensor): + The encoder feature that generates ys (n_batch, xlen, n_feat). + + Returns: + tuple[torch.Tensor, List[Any]]: Tuple of + batchfied scores for next token with shape of `(n_batch, n_vocab)` + and next state list for ys. + + """ + # merge states + n_batch = len(ys) + n_layers = len(self.decoders) + if states[0] is None: + batch_state = None + else: + # transpose state of [batch, layer] into [layer, batch] + batch_state = [ + torch.stack([states[b][i] for b in range(n_batch)]) for i in range(n_layers) + ] + + # batch decoding + ys_mask = subsequent_mask(ys.size(-1), device=xs.device).unsqueeze(0) + logp, states = self.forward_one_step(ys, ys_mask, xs, cache=batch_state) + + # transpose state of [layer, batch] into [batch, layer] + state_list = [[states[i][b] for i in range(n_layers)] for b in range(n_batch)] + return logp, state_list + + +@tables.register("decoder_classes", "TransformerDecoder") +class TransformerDecoder(BaseTransformerDecoder): + def __init__( + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + ): + super().__init__( + vocab_size=vocab_size, + encoder_output_size=encoder_output_size, + dropout_rate=dropout_rate, + positional_dropout_rate=positional_dropout_rate, + input_layer=input_layer, + use_output_layer=use_output_layer, + pos_enc_class=pos_enc_class, + normalize_before=normalize_before, + ) + + attention_dim = encoder_output_size + self.decoders = repeat( + num_blocks, + lambda lnum: DecoderLayer( + attention_dim, + MultiHeadedAttention(attention_heads, attention_dim, self_attention_dropout_rate), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), + PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), + dropout_rate, + normalize_before, + concat_after, + ), + ) + + +@tables.register("decoder_classes", "LightweightConvolutionTransformerDecoder") +class LightweightConvolutionTransformerDecoder(BaseTransformerDecoder): + def __init__( + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + conv_wshare: int = 4, + conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), + conv_usebias: int = False, + ): + if len(conv_kernel_length) != num_blocks: + raise ValueError( + "conv_kernel_length must have equal number of values to num_blocks: " + f"{len(conv_kernel_length)} != {num_blocks}" + ) + super().__init__( + vocab_size=vocab_size, + encoder_output_size=encoder_output_size, + dropout_rate=dropout_rate, + positional_dropout_rate=positional_dropout_rate, + input_layer=input_layer, + use_output_layer=use_output_layer, + pos_enc_class=pos_enc_class, + normalize_before=normalize_before, + ) + + attention_dim = encoder_output_size + self.decoders = repeat( + num_blocks, + lambda lnum: DecoderLayer( + attention_dim, + LightweightConvolution( + wshare=conv_wshare, + n_feat=attention_dim, + dropout_rate=self_attention_dropout_rate, + kernel_size=conv_kernel_length[lnum], + use_kernel_mask=True, + use_bias=conv_usebias, + ), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), + PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), + dropout_rate, + normalize_before, + concat_after, + ), + ) + + +@tables.register("decoder_classes", "LightweightConvolution2DTransformerDecoder") +class LightweightConvolution2DTransformerDecoder(BaseTransformerDecoder): + def __init__( + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + conv_wshare: int = 4, + conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), + conv_usebias: int = False, + ): + if len(conv_kernel_length) != num_blocks: + raise ValueError( + "conv_kernel_length must have equal number of values to num_blocks: " + f"{len(conv_kernel_length)} != {num_blocks}" + ) + super().__init__( + vocab_size=vocab_size, + encoder_output_size=encoder_output_size, + dropout_rate=dropout_rate, + positional_dropout_rate=positional_dropout_rate, + input_layer=input_layer, + use_output_layer=use_output_layer, + pos_enc_class=pos_enc_class, + normalize_before=normalize_before, + ) + + attention_dim = encoder_output_size + self.decoders = repeat( + num_blocks, + lambda lnum: DecoderLayer( + attention_dim, + LightweightConvolution2D( + wshare=conv_wshare, + n_feat=attention_dim, + dropout_rate=self_attention_dropout_rate, + kernel_size=conv_kernel_length[lnum], + use_kernel_mask=True, + use_bias=conv_usebias, + ), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), + PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), + dropout_rate, + normalize_before, + concat_after, + ), + ) + + +@tables.register("decoder_classes", "DynamicConvolutionTransformerDecoder") +class DynamicConvolutionTransformerDecoder(BaseTransformerDecoder): + def __init__( + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + conv_wshare: int = 4, + conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), + conv_usebias: int = False, + ): + if len(conv_kernel_length) != num_blocks: + raise ValueError( + "conv_kernel_length must have equal number of values to num_blocks: " + f"{len(conv_kernel_length)} != {num_blocks}" + ) + super().__init__( + vocab_size=vocab_size, + encoder_output_size=encoder_output_size, + dropout_rate=dropout_rate, + positional_dropout_rate=positional_dropout_rate, + input_layer=input_layer, + use_output_layer=use_output_layer, + pos_enc_class=pos_enc_class, + normalize_before=normalize_before, + ) + attention_dim = encoder_output_size + + self.decoders = repeat( + num_blocks, + lambda lnum: DecoderLayer( + attention_dim, + DynamicConvolution( + wshare=conv_wshare, + n_feat=attention_dim, + dropout_rate=self_attention_dropout_rate, + kernel_size=conv_kernel_length[lnum], + use_kernel_mask=True, + use_bias=conv_usebias, + ), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), + PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), + dropout_rate, + normalize_before, + concat_after, + ), + ) + + +@tables.register("decoder_classes", "DynamicConvolution2DTransformerDecoder") +class DynamicConvolution2DTransformerDecoder(BaseTransformerDecoder): + def __init__( + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + conv_wshare: int = 4, + conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), + conv_usebias: int = False, + ): + if len(conv_kernel_length) != num_blocks: + raise ValueError( + "conv_kernel_length must have equal number of values to num_blocks: " + f"{len(conv_kernel_length)} != {num_blocks}" + ) + super().__init__( + vocab_size=vocab_size, + encoder_output_size=encoder_output_size, + dropout_rate=dropout_rate, + positional_dropout_rate=positional_dropout_rate, + input_layer=input_layer, + use_output_layer=use_output_layer, + pos_enc_class=pos_enc_class, + normalize_before=normalize_before, + ) + attention_dim = encoder_output_size + + self.decoders = repeat( + num_blocks, + lambda lnum: DecoderLayer( + attention_dim, + DynamicConvolution2D( + wshare=conv_wshare, + n_feat=attention_dim, + dropout_rate=self_attention_dropout_rate, + kernel_size=conv_kernel_length[lnum], + use_kernel_mask=True, + use_bias=conv_usebias, + ), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), + PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), + dropout_rate, + normalize_before, + concat_after, + ), + ) diff --git a/funasr/models/transformer/embedding.py b/funasr/models/transformer/embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..7c86dab2ec77794e9012b82a71dc761cf52b704a --- /dev/null +++ b/funasr/models/transformer/embedding.py @@ -0,0 +1,533 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +# Copyright 2019 Shigeki Karita +# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +"""Positional Encoding Module.""" + +import math +import torch +import torch.nn.functional as F +from torch import einsum + + +def _pre_hook( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, +): + """Perform pre-hook in load_state_dict for backward compatibility. + + Note: + We saved self.pe until v.0.5.2 but we have omitted it later. + Therefore, we remove the item "pe" from `state_dict` for backward compatibility. + + """ + k = prefix + "pe" + if k in state_dict: + state_dict.pop(k) + + +class PositionalEncoding(torch.nn.Module): + """Positional encoding. + + Args: + d_model (int): Embedding dimension. + dropout_rate (float): Dropout rate. + max_len (int): Maximum input length. + reverse (bool): Whether to reverse the input position. Only for + the class LegacyRelPositionalEncoding. We remove it in the current + class RelPositionalEncoding. + """ + + def __init__(self, d_model, dropout_rate, max_len=5000, reverse=False): + """Construct an PositionalEncoding object.""" + super(PositionalEncoding, self).__init__() + self.d_model = d_model + self.reverse = reverse + self.xscale = math.sqrt(self.d_model) + self.dropout = torch.nn.Dropout(p=dropout_rate) + self.pe = None + self.extend_pe(torch.tensor(0.0).expand(1, max_len)) + self._register_load_state_dict_pre_hook(_pre_hook) + + def extend_pe(self, x): + """Reset the positional encodings.""" + if self.pe is not None: + if self.pe.size(1) >= x.size(1): + if self.pe.dtype != x.dtype or self.pe.device != x.device: + self.pe = self.pe.to(dtype=x.dtype, device=x.device) + return + pe = torch.zeros(x.size(1), self.d_model) + if self.reverse: + position = torch.arange(x.size(1) - 1, -1, -1.0, dtype=torch.float32).unsqueeze(1) + else: + position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1) + div_term = torch.exp( + torch.arange(0, self.d_model, 2, dtype=torch.float32) + * -(math.log(10000.0) / self.d_model) + ) + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + pe = pe.unsqueeze(0) + self.pe = pe.to(device=x.device, dtype=x.dtype) + + def forward(self, x: torch.Tensor): + """Add positional encoding. + + Args: + x (torch.Tensor): Input tensor (batch, time, `*`). + + Returns: + torch.Tensor: Encoded tensor (batch, time, `*`). + """ + self.extend_pe(x) + x = x * self.xscale + self.pe[:, : x.size(1)] + return self.dropout(x) + + +class ScaledPositionalEncoding(PositionalEncoding): + """Scaled positional encoding module. + + See Sec. 3.2 https://arxiv.org/abs/1809.08895 + + Args: + d_model (int): Embedding dimension. + dropout_rate (float): Dropout rate. + max_len (int): Maximum input length. + + """ + + def __init__(self, d_model, dropout_rate, max_len=5000): + """Initialize class.""" + super().__init__(d_model=d_model, dropout_rate=dropout_rate, max_len=max_len) + self.alpha = torch.nn.Parameter(torch.tensor(1.0)) + + def reset_parameters(self): + """Reset parameters.""" + self.alpha.data = torch.tensor(1.0) + + def forward(self, x): + """Add positional encoding. + + Args: + x (torch.Tensor): Input tensor (batch, time, `*`). + + Returns: + torch.Tensor: Encoded tensor (batch, time, `*`). + + """ + self.extend_pe(x) + x = x + self.alpha * self.pe[:, : x.size(1)] + return self.dropout(x) + + +class LearnableFourierPosEnc(torch.nn.Module): + """Learnable Fourier Features for Positional Encoding. + + See https://arxiv.org/pdf/2106.02795.pdf + + Args: + d_model (int): Embedding dimension. + dropout_rate (float): Dropout rate. + max_len (int): Maximum input length. + gamma (float): init parameter for the positional kernel variance + see https://arxiv.org/pdf/2106.02795.pdf. + apply_scaling (bool): Whether to scale the input before adding the pos encoding. + hidden_dim (int): if not None, we modulate the pos encodings with + an MLP whose hidden layer has hidden_dim neurons. + """ + + def __init__( + self, + d_model, + dropout_rate=0.0, + max_len=5000, + gamma=1.0, + apply_scaling=False, + hidden_dim=None, + ): + """Initialize class.""" + super(LearnableFourierPosEnc, self).__init__() + + self.d_model = d_model + + if apply_scaling: + self.xscale = math.sqrt(self.d_model) + else: + self.xscale = 1.0 + + self.dropout = torch.nn.Dropout(dropout_rate) + self.max_len = max_len + + self.gamma = gamma + if self.gamma is None: + self.gamma = self.d_model // 2 + + assert d_model % 2 == 0, "d_model should be divisible by two in order to use this layer." + self.w_r = torch.nn.Parameter(torch.empty(1, d_model // 2)) + self._reset() # init the weights + + self.hidden_dim = hidden_dim + if self.hidden_dim is not None: + self.mlp = torch.nn.Sequential( + torch.nn.Linear(d_model, hidden_dim), + torch.nn.GELU(), + torch.nn.Linear(hidden_dim, d_model), + ) + + def _reset(self): + self.w_r.data = torch.normal(0, (1 / math.sqrt(self.gamma)), (1, self.d_model // 2)) + + def extend_pe(self, x): + """Reset the positional encodings.""" + position_v = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1).to(x) + + cosine = torch.cos(torch.matmul(position_v, self.w_r)) + sine = torch.sin(torch.matmul(position_v, self.w_r)) + pos_enc = torch.cat((cosine, sine), -1) + pos_enc /= math.sqrt(self.d_model) + + if self.hidden_dim is None: + return pos_enc.unsqueeze(0) + else: + return self.mlp(pos_enc.unsqueeze(0)) + + def forward(self, x: torch.Tensor): + """Add positional encoding. + + Args: + x (torch.Tensor): Input tensor (batch, time, `*`). + + Returns: + torch.Tensor: Encoded tensor (batch, time, `*`). + """ + pe = self.extend_pe(x) + x = x * self.xscale + pe + return self.dropout(x) + + +class LegacyRelPositionalEncoding(PositionalEncoding): + """Relative positional encoding module (old version). + + Details can be found in https://github.com/espnet/espnet/pull/2816. + + See : Appendix B in https://arxiv.org/abs/1901.02860 + + Args: + d_model (int): Embedding dimension. + dropout_rate (float): Dropout rate. + max_len (int): Maximum input length. + + """ + + def __init__(self, d_model, dropout_rate, max_len=5000): + """Initialize class.""" + super().__init__( + d_model=d_model, + dropout_rate=dropout_rate, + max_len=max_len, + reverse=True, + ) + + def forward(self, x): + """Compute positional encoding. + + Args: + x (torch.Tensor): Input tensor (batch, time, `*`). + + Returns: + torch.Tensor: Encoded tensor (batch, time, `*`). + torch.Tensor: Positional embedding tensor (1, time, `*`). + + """ + self.extend_pe(x) + x = x * self.xscale + pos_emb = self.pe[:, : x.size(1)] + return self.dropout(x), self.dropout(pos_emb) + + +class RelPositionalEncoding(torch.nn.Module): + """Relative positional encoding module (new implementation). + + Details can be found in https://github.com/espnet/espnet/pull/2816. + + See : Appendix B in https://arxiv.org/abs/1901.02860 + + Args: + d_model (int): Embedding dimension. + dropout_rate (float): Dropout rate. + max_len (int): Maximum input length. + + """ + + def __init__(self, d_model, dropout_rate, max_len=5000): + """Construct an PositionalEncoding object.""" + super(RelPositionalEncoding, self).__init__() + self.d_model = d_model + self.xscale = math.sqrt(self.d_model) + self.dropout = torch.nn.Dropout(p=dropout_rate) + self.pe = None + self.extend_pe(torch.tensor(0.0).expand(1, max_len)) + + def extend_pe(self, x): + """Reset the positional encodings.""" + if self.pe is not None: + # self.pe contains both positive and negative parts + # the length of self.pe is 2 * input_len - 1 + if self.pe.size(1) >= x.size(1) * 2 - 1: + if self.pe.dtype != x.dtype or self.pe.device != x.device: + self.pe = self.pe.to(dtype=x.dtype, device=x.device) + return + # Suppose `i` means to the position of query vecotr and `j` means the + # position of key vector. We use position relative positions when keys + # are to the left (i>j) and negative relative positions otherwise (i= length: + if self.pe.dtype != dtype or self.pe.device != device: + self.pe = self.pe.to(dtype=dtype, device=device) + return + pe = torch.zeros(length, self.d_model) + position = torch.arange(0, length, dtype=torch.float32).unsqueeze(1) + div_term = torch.exp( + torch.arange(0, self.d_model, 2, dtype=torch.float32) + * -(math.log(10000.0) / self.d_model) + ) + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + pe = pe.unsqueeze(0) + self.pe = pe.to(device=device, dtype=dtype) + + def forward(self, x: torch.Tensor, start_idx: int = 0): + """Add positional encoding. + + Args: + x (torch.Tensor): Input tensor (batch, time, `*`). + + Returns: + torch.Tensor: Encoded tensor (batch, time, `*`). + + """ + self.extend_pe(x.size(1) + start_idx, x.device, x.dtype) + x = x * self.xscale + self.pe[:, start_idx : start_idx + x.size(1)] + return self.dropout(x) + + +class SinusoidalPositionEncoder(torch.nn.Module): + """ """ + + def __int__(self, d_model=80, dropout_rate=0.1): + pass + + def encode( + self, positions: torch.Tensor = None, depth: int = None, dtype: torch.dtype = torch.float32 + ): + batch_size = positions.size(0) + positions = positions.type(dtype) + device = positions.device + log_timescale_increment = torch.log(torch.tensor([10000], dtype=dtype, device=device)) / ( + depth / 2 - 1 + ) + inv_timescales = torch.exp( + torch.arange(depth / 2, device=device).type(dtype) * (-log_timescale_increment) + ) + inv_timescales = torch.reshape(inv_timescales, [batch_size, -1]) + scaled_time = torch.reshape(positions, [1, -1, 1]) * torch.reshape( + inv_timescales, [1, 1, -1] + ) + encoding = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=2) + return encoding.type(dtype) + + def forward(self, x): + batch_size, timesteps, input_dim = x.size() + positions = torch.arange(1, timesteps + 1, device=x.device)[None, :] + position_encoding = self.encode(positions, input_dim, x.dtype).to(x.device) + + return x + position_encoding + + +class StreamSinusoidalPositionEncoder(torch.nn.Module): + """ """ + + def __int__(self, d_model=80, dropout_rate=0.1): + pass + + def encode( + self, positions: torch.Tensor = None, depth: int = None, dtype: torch.dtype = torch.float32 + ): + batch_size = positions.size(0) + positions = positions.type(dtype) + log_timescale_increment = torch.log(torch.tensor([10000], dtype=dtype)) / (depth / 2 - 1) + inv_timescales = torch.exp(torch.arange(depth / 2).type(dtype) * (-log_timescale_increment)) + inv_timescales = torch.reshape(inv_timescales, [batch_size, -1]) + scaled_time = torch.reshape(positions, [1, -1, 1]) * torch.reshape( + inv_timescales, [1, 1, -1] + ) + encoding = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=2) + return encoding.type(dtype) + + def forward(self, x, cache=None): + batch_size, timesteps, input_dim = x.size() + start_idx = 0 + if cache is not None: + start_idx = cache["start_idx"] + cache["start_idx"] += timesteps + positions = torch.arange(1, timesteps + start_idx + 1)[None, :] + position_encoding = self.encode(positions, input_dim, x.dtype).to(x.device) + return x + position_encoding[:, start_idx : start_idx + timesteps] + + +class StreamingRelPositionalEncoding(torch.nn.Module): + """Relative positional encoding. + Args: + size: Module size. + max_len: Maximum input length. + dropout_rate: Dropout rate. + """ + + def __init__(self, size: int, dropout_rate: float = 0.0, max_len: int = 5000) -> None: + """Construct a RelativePositionalEncoding object.""" + super().__init__() + + self.size = size + + self.pe = None + self.dropout = torch.nn.Dropout(p=dropout_rate) + + self.extend_pe(torch.tensor(0.0).expand(1, max_len)) + self._register_load_state_dict_pre_hook(_pre_hook) + + def extend_pe(self, x: torch.Tensor, left_context: int = 0) -> None: + """Reset positional encoding. + Args: + x: Input sequences. (B, T, ?) + left_context: Number of frames in left context. + """ + time1 = x.size(1) + left_context + + if self.pe is not None: + if self.pe.size(1) >= time1 * 2 - 1: + if self.pe.dtype != x.dtype or self.pe.device != x.device: + self.pe = self.pe.to(device=x.device, dtype=x.dtype) + return + + pe_positive = torch.zeros(time1, self.size) + pe_negative = torch.zeros(time1, self.size) + + position = torch.arange(0, time1, dtype=torch.float32).unsqueeze(1) + div_term = torch.exp( + torch.arange(0, self.size, 2, dtype=torch.float32) * -(math.log(10000.0) / self.size) + ) + + pe_positive[:, 0::2] = torch.sin(position * div_term) + pe_positive[:, 1::2] = torch.cos(position * div_term) + pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0) + + pe_negative[:, 0::2] = torch.sin(-1 * position * div_term) + pe_negative[:, 1::2] = torch.cos(-1 * position * div_term) + pe_negative = pe_negative[1:].unsqueeze(0) + + self.pe = torch.cat([pe_positive, pe_negative], dim=1).to(dtype=x.dtype, device=x.device) + + def forward(self, x: torch.Tensor, left_context: int = 0) -> torch.Tensor: + """Compute positional encoding. + Args: + x: Input sequences. (B, T, ?) + left_context: Number of frames in left context. + Returns: + pos_enc: Positional embedding sequences. (B, 2 * (T - 1), ?) + """ + self.extend_pe(x, left_context=left_context) + + time1 = x.size(1) + left_context + + pos_enc = self.pe[:, self.pe.size(1) // 2 - time1 + 1 : self.pe.size(1) // 2 + x.size(1)] + pos_enc = self.dropout(pos_enc) + + return pos_enc + + +class ScaledSinuEmbedding(torch.nn.Module): + def __init__(self, dim): + super().__init__() + self.scale = torch.nn.Parameter( + torch.ones( + 1, + ) + ) + inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim)) + self.register_buffer("inv_freq", inv_freq) + + def forward(self, x): + n, device = x.shape[1], x.device + t = torch.arange(n, device=device).type_as(self.inv_freq) + sinu = einsum("i , j -> i j", t, self.inv_freq) + emb = torch.cat((sinu.sin(), sinu.cos()), dim=-1) + return emb * self.scale diff --git a/funasr/models/transformer/encoder.py b/funasr/models/transformer/encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..987924f0cf66bc4a3cb9334aa0dd7db4384e37f6 --- /dev/null +++ b/funasr/models/transformer/encoder.py @@ -0,0 +1,329 @@ +# Copyright 2019 Shigeki Karita +# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +"""Transformer encoder definition.""" + +from typing import List +from typing import Optional +from typing import Tuple + +import torch +from torch import nn +import logging + +from funasr.models.transformer.attention import MultiHeadedAttention +from funasr.models.transformer.embedding import PositionalEncoding +from funasr.models.transformer.layer_norm import LayerNorm +from funasr.models.transformer.utils.multi_layer_conv import Conv1dLinear +from funasr.models.transformer.utils.multi_layer_conv import MultiLayeredConv1d +from funasr.models.transformer.utils.nets_utils import make_pad_mask +from funasr.models.transformer.positionwise_feed_forward import PositionwiseFeedForward +from funasr.models.transformer.utils.repeat import repeat +from funasr.models.ctc.ctc import CTC + +from funasr.models.transformer.utils.subsampling import Conv2dSubsampling +from funasr.models.transformer.utils.subsampling import Conv2dSubsampling2 +from funasr.models.transformer.utils.subsampling import Conv2dSubsampling6 +from funasr.models.transformer.utils.subsampling import Conv2dSubsampling8 +from funasr.models.transformer.utils.subsampling import TooShortUttError +from funasr.models.transformer.utils.subsampling import check_short_utt + +from funasr.register import tables + + +class EncoderLayer(nn.Module): + """Encoder layer module. + + Args: + size (int): Input dimension. + self_attn (torch.nn.Module): Self-attention module instance. + `MultiHeadedAttention` or `RelPositionMultiHeadedAttention` instance + can be used as the argument. + feed_forward (torch.nn.Module): Feed-forward module instance. + `PositionwiseFeedForward`, `MultiLayeredConv1d`, or `Conv1dLinear` instance + can be used as the argument. + dropout_rate (float): Dropout rate. + normalize_before (bool): Whether to use layer_norm before the first block. + concat_after (bool): Whether to concat attention layer's input and output. + if True, additional linear will be applied. + i.e. x -> x + linear(concat(x, att(x))) + if False, no additional linear will be applied. i.e. x -> x + att(x) + stochastic_depth_rate (float): Proability to skip this layer. + During training, the layer may skip residual computation and return input + as-is with given probability. + """ + + def __init__( + self, + size, + self_attn, + feed_forward, + dropout_rate, + normalize_before=True, + concat_after=False, + stochastic_depth_rate=0.0, + ): + """Construct an EncoderLayer object.""" + super().__init__() + self.self_attn = self_attn + self.feed_forward = feed_forward + self.norm1 = LayerNorm(size) + self.norm2 = LayerNorm(size) + self.dropout = nn.Dropout(dropout_rate) + self.size = size + self.normalize_before = normalize_before + self.concat_after = concat_after + if self.concat_after: + self.concat_linear = nn.Linear(size + size, size) + self.stochastic_depth_rate = stochastic_depth_rate + + def forward(self, x, mask, cache=None): + """Compute encoded features. + + Args: + x_input (torch.Tensor): Input tensor (#batch, time, size). + mask (torch.Tensor): Mask tensor for the input (#batch, time). + cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size). + + Returns: + torch.Tensor: Output tensor (#batch, time, size). + torch.Tensor: Mask tensor (#batch, time). + + """ + skip_layer = False + # with stochastic depth, residual connection `x + f(x)` becomes + # `x <- x + 1 / (1 - p) * f(x)` at training time. + stoch_layer_coeff = 1.0 + if self.training and self.stochastic_depth_rate > 0: + skip_layer = torch.rand(1).item() < self.stochastic_depth_rate + stoch_layer_coeff = 1.0 / (1 - self.stochastic_depth_rate) + + if skip_layer: + if cache is not None: + x = torch.cat([cache, x], dim=1) + return x, mask + + residual = x + if self.normalize_before: + x = self.norm1(x) + + if cache is None: + x_q = x + else: + assert cache.shape == (x.shape[0], x.shape[1] - 1, self.size) + x_q = x[:, -1:, :] + residual = residual[:, -1:, :] + mask = None if mask is None else mask[:, -1:, :] + + if self.concat_after: + x_concat = torch.cat((x, self.self_attn(x_q, x, x, mask)), dim=-1) + x = residual + stoch_layer_coeff * self.concat_linear(x_concat) + else: + x = residual + stoch_layer_coeff * self.dropout(self.self_attn(x_q, x, x, mask)) + if not self.normalize_before: + x = self.norm1(x) + + residual = x + if self.normalize_before: + x = self.norm2(x) + x = residual + stoch_layer_coeff * self.dropout(self.feed_forward(x)) + if not self.normalize_before: + x = self.norm2(x) + + if cache is not None: + x = torch.cat([cache, x], dim=1) + + return x, mask + + +@tables.register("encoder_classes", "TransformerEncoder") +class TransformerEncoder(nn.Module): + """Transformer encoder module. + + Args: + input_size: input dim + output_size: dimension of attention + attention_heads: the number of heads of multi head attention + linear_units: the number of units of position-wise feed forward + num_blocks: the number of decoder blocks + dropout_rate: dropout rate + attention_dropout_rate: dropout rate in attention + positional_dropout_rate: dropout rate after adding positional encoding + input_layer: input layer type + pos_enc_class: PositionalEncoding or ScaledPositionalEncoding + normalize_before: whether to use layer_norm before the first block + concat_after: whether to concat attention layer's input and output + if True, additional linear will be applied. + i.e. x -> x + linear(concat(x, att(x))) + if False, no additional linear will be applied. + i.e. x -> x + att(x) + positionwise_layer_type: linear of conv1d + positionwise_conv_kernel_size: kernel size of positionwise conv1d layer + padding_idx: padding_idx for input_layer=embed + """ + + def __init__( + self, + input_size: int, + output_size: int = 256, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + attention_dropout_rate: float = 0.0, + input_layer: Optional[str] = "conv2d", + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + positionwise_layer_type: str = "linear", + positionwise_conv_kernel_size: int = 1, + padding_idx: int = -1, + interctc_layer_idx: List[int] = [], + interctc_use_conditioning: bool = False, + ): + super().__init__() + self._output_size = output_size + + if input_layer == "linear": + self.embed = torch.nn.Sequential( + torch.nn.Linear(input_size, output_size), + torch.nn.LayerNorm(output_size), + torch.nn.Dropout(dropout_rate), + torch.nn.ReLU(), + pos_enc_class(output_size, positional_dropout_rate), + ) + elif input_layer == "conv2d": + self.embed = Conv2dSubsampling(input_size, output_size, dropout_rate) + elif input_layer == "conv2d2": + self.embed = Conv2dSubsampling2(input_size, output_size, dropout_rate) + elif input_layer == "conv2d6": + self.embed = Conv2dSubsampling6(input_size, output_size, dropout_rate) + elif input_layer == "conv2d8": + self.embed = Conv2dSubsampling8(input_size, output_size, dropout_rate) + elif input_layer == "embed": + self.embed = torch.nn.Sequential( + torch.nn.Embedding(input_size, output_size, padding_idx=padding_idx), + pos_enc_class(output_size, positional_dropout_rate), + ) + elif input_layer is None: + if input_size == output_size: + self.embed = None + else: + self.embed = torch.nn.Linear(input_size, output_size) + else: + raise ValueError("unknown input_layer: " + input_layer) + self.normalize_before = normalize_before + if positionwise_layer_type == "linear": + positionwise_layer = PositionwiseFeedForward + positionwise_layer_args = ( + output_size, + linear_units, + dropout_rate, + ) + elif positionwise_layer_type == "conv1d": + positionwise_layer = MultiLayeredConv1d + positionwise_layer_args = ( + output_size, + linear_units, + positionwise_conv_kernel_size, + dropout_rate, + ) + elif positionwise_layer_type == "conv1d-linear": + positionwise_layer = Conv1dLinear + positionwise_layer_args = ( + output_size, + linear_units, + positionwise_conv_kernel_size, + dropout_rate, + ) + else: + raise NotImplementedError("Support only linear or conv1d.") + self.encoders = repeat( + num_blocks, + lambda lnum: EncoderLayer( + output_size, + MultiHeadedAttention(attention_heads, output_size, attention_dropout_rate), + positionwise_layer(*positionwise_layer_args), + dropout_rate, + normalize_before, + concat_after, + ), + ) + if self.normalize_before: + self.after_norm = LayerNorm(output_size) + + self.interctc_layer_idx = interctc_layer_idx + if len(interctc_layer_idx) > 0: + assert 0 < min(interctc_layer_idx) and max(interctc_layer_idx) < num_blocks + self.interctc_use_conditioning = interctc_use_conditioning + self.conditioning_layer = None + + def output_size(self) -> int: + return self._output_size + + def forward( + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + prev_states: torch.Tensor = None, + ctc: CTC = None, + ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Embed positions in tensor. + + Args: + xs_pad: input tensor (B, L, D) + ilens: input length (B) + prev_states: Not to be used now. + Returns: + position embedded tensor and mask + """ + masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device) + + if self.embed is None: + xs_pad = xs_pad + elif ( + isinstance(self.embed, Conv2dSubsampling) + or isinstance(self.embed, Conv2dSubsampling2) + or isinstance(self.embed, Conv2dSubsampling6) + or isinstance(self.embed, Conv2dSubsampling8) + ): + short_status, limit_size = check_short_utt(self.embed, xs_pad.size(1)) + if short_status: + raise TooShortUttError( + f"has {xs_pad.size(1)} frames and is too short for subsampling " + + f"(it needs more than {limit_size} frames), return empty results", + xs_pad.size(1), + limit_size, + ) + xs_pad, masks = self.embed(xs_pad, masks) + else: + xs_pad = self.embed(xs_pad) + + intermediate_outs = [] + if len(self.interctc_layer_idx) == 0: + xs_pad, masks = self.encoders(xs_pad, masks) + else: + for layer_idx, encoder_layer in enumerate(self.encoders): + xs_pad, masks = encoder_layer(xs_pad, masks) + + if layer_idx + 1 in self.interctc_layer_idx: + encoder_out = xs_pad + + # intermediate outputs are also normalized + if self.normalize_before: + encoder_out = self.after_norm(encoder_out) + + intermediate_outs.append((layer_idx + 1, encoder_out)) + + if self.interctc_use_conditioning: + ctc_out = ctc.softmax(encoder_out) + xs_pad = xs_pad + self.conditioning_layer(ctc_out) + + if self.normalize_before: + xs_pad = self.after_norm(xs_pad) + + olens = masks.squeeze(1).sum(1) + if len(intermediate_outs) > 0: + return (xs_pad, intermediate_outs), olens, None + return xs_pad, olens, None diff --git a/funasr/models/transformer/layer_norm.py b/funasr/models/transformer/layer_norm.py new file mode 100644 index 0000000000000000000000000000000000000000..ce4f8877f861399f517bc9d46197aba5780bbf03 --- /dev/null +++ b/funasr/models/transformer/layer_norm.py @@ -0,0 +1,166 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +# Copyright 2019 Shigeki Karita +# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +"""Layer normalization module.""" + +import torch +import torch.nn as nn + + +class LayerNorm(torch.nn.LayerNorm): + """Layer normalization module. + + Args: + nout (int): Output dim size. + dim (int): Dimension to be normalized. + + """ + + def __init__(self, nout, dim=-1): + """Construct an LayerNorm object.""" + super(LayerNorm, self).__init__(nout, eps=1e-12) + self.dim = dim + + def forward(self, x): + """Apply layer normalization. + + Args: + x (torch.Tensor): Input tensor. + + Returns: + torch.Tensor: Normalized tensor. + + """ + if self.dim == -1: + return super(LayerNorm, self).forward(x) + return super(LayerNorm, self).forward(x.transpose(self.dim, -1)).transpose(self.dim, -1) + + +class GlobalLayerNorm(nn.Module): + """Calculate Global Layer Normalization. + + Arguments + --------- + dim : (int or list or torch.Size) + Input shape from an expected input of size. + eps : float + A value added to the denominator for numerical stability. + elementwise_affine : bool + A boolean value that when set to True, + this module has learnable per-element affine parameters + initialized to ones (for weights) and zeros (for biases). + + Example + ------- + >>> x = torch.randn(5, 10, 20) + >>> GLN = GlobalLayerNorm(10, 3) + >>> x_norm = GLN(x) + """ + + def __init__(self, dim, shape, eps=1e-8, elementwise_affine=True): + super(GlobalLayerNorm, self).__init__() + self.dim = dim + self.eps = eps + self.elementwise_affine = elementwise_affine + + if self.elementwise_affine: + if shape == 3: + self.weight = nn.Parameter(torch.ones(self.dim, 1)) + self.bias = nn.Parameter(torch.zeros(self.dim, 1)) + if shape == 4: + self.weight = nn.Parameter(torch.ones(self.dim, 1, 1)) + self.bias = nn.Parameter(torch.zeros(self.dim, 1, 1)) + else: + self.register_parameter("weight", None) + self.register_parameter("bias", None) + + def forward(self, x): + """Returns the normalized tensor. + + Arguments + --------- + x : torch.Tensor + Tensor of size [N, C, K, S] or [N, C, L]. + """ + # x = N x C x K x S or N x C x L + # N x 1 x 1 + # cln: mean,var N x 1 x K x S + # gln: mean,var N x 1 x 1 + if x.dim() == 3: + mean = torch.mean(x, (1, 2), keepdim=True) + var = torch.mean((x - mean) ** 2, (1, 2), keepdim=True) + if self.elementwise_affine: + x = self.weight * (x - mean) / torch.sqrt(var + self.eps) + self.bias + else: + x = (x - mean) / torch.sqrt(var + self.eps) + + if x.dim() == 4: + mean = torch.mean(x, (1, 2, 3), keepdim=True) + var = torch.mean((x - mean) ** 2, (1, 2, 3), keepdim=True) + if self.elementwise_affine: + x = self.weight * (x - mean) / torch.sqrt(var + self.eps) + self.bias + else: + x = (x - mean) / torch.sqrt(var + self.eps) + return x + + +class CumulativeLayerNorm(nn.LayerNorm): + """Calculate Cumulative Layer Normalization. + + Arguments + --------- + dim : int + Dimension that you want to normalize. + elementwise_affine : True + Learnable per-element affine parameters. + + Example + ------- + >>> x = torch.randn(5, 10, 20) + >>> CLN = CumulativeLayerNorm(10) + >>> x_norm = CLN(x) + """ + + def __init__(self, dim, elementwise_affine=True): + super(CumulativeLayerNorm, self).__init__( + dim, elementwise_affine=elementwise_affine, eps=1e-8 + ) + + def forward(self, x): + """Returns the normalized tensor. + + Arguments + --------- + x : torch.Tensor + Tensor size [N, C, K, S] or [N, C, L] + """ + # x: N x C x K x S or N x C x L + # N x K x S x C + if x.dim() == 4: + x = x.permute(0, 2, 3, 1).contiguous() + # N x K x S x C == only channel norm + x = super().forward(x) + # N x C x K x S + x = x.permute(0, 3, 1, 2).contiguous() + if x.dim() == 3: + x = torch.transpose(x, 1, 2) + # N x L x C == only channel norm + x = super().forward(x) + # N x C x L + x = torch.transpose(x, 1, 2) + return x + + +class ScaleNorm(nn.Module): + def __init__(self, dim, eps=1e-5): + super().__init__() + self.scale = dim**-0.5 + self.eps = eps + self.g = nn.Parameter(torch.ones(1)) + + def forward(self, x): + norm = torch.norm(x, dim=-1, keepdim=True) * self.scale + return x / norm.clamp(min=self.eps) * self.g diff --git a/funasr/models/transformer/model.py b/funasr/models/transformer/model.py new file mode 100644 index 0000000000000000000000000000000000000000..adfd52520bb954585229c301ca3f803117f0fca8 --- /dev/null +++ b/funasr/models/transformer/model.py @@ -0,0 +1,459 @@ +import logging +from typing import Union, Dict, List, Tuple, Optional + +import time +import torch +import torch.nn as nn +from torch.cuda.amp import autocast + +from funasr.losses.label_smoothing_loss import LabelSmoothingLoss +from funasr.models.ctc.ctc import CTC +from funasr.models.transformer.utils.add_sos_eos import add_sos_eos +from funasr.metrics.compute_acc import th_accuracy + +# from funasr.models.e2e_asr_common import ErrorCalculator +from funasr.train_utils.device_funcs import force_gatherable +from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank +from funasr.utils import postprocess_utils +from funasr.utils.datadir_writer import DatadirWriter +from funasr.register import tables + + +@tables.register("model_classes", "Transformer") +class Transformer(nn.Module): + """CTC-attention hybrid Encoder-Decoder model""" + + def __init__( + self, + specaug: str = None, + specaug_conf: dict = None, + normalize: str = None, + normalize_conf: dict = None, + encoder: str = None, + encoder_conf: dict = None, + decoder: str = None, + decoder_conf: dict = None, + ctc: str = None, + ctc_conf: dict = None, + ctc_weight: float = 0.5, + interctc_weight: float = 0.0, + input_size: int = 80, + vocab_size: int = -1, + ignore_id: int = -1, + blank_id: int = 0, + sos: int = 1, + eos: int = 2, + lsm_weight: float = 0.0, + length_normalized_loss: bool = False, + report_cer: bool = True, + report_wer: bool = True, + sym_space: str = "", + sym_blank: str = "", + # extract_feats_in_collect_stats: bool = True, + share_embedding: bool = False, + # preencoder: Optional[AbsPreEncoder] = None, + # postencoder: Optional[AbsPostEncoder] = None, + **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() + if decoder is not None: + decoder_class = tables.decoder_classes.get(decoder) + decoder = decoder_class( + vocab_size=vocab_size, + encoder_output_size=encoder_output_size, + **decoder_conf, + ) + if ctc_weight > 0.0: + + if ctc_conf is None: + ctc_conf = {} + + ctc = CTC(odim=vocab_size, encoder_output_size=encoder_output_size, **ctc_conf) + + self.blank_id = blank_id + self.sos = sos if sos is not None else vocab_size - 1 + self.eos = eos if eos is not None else vocab_size - 1 + self.vocab_size = vocab_size + self.ignore_id = ignore_id + self.ctc_weight = ctc_weight + self.specaug = specaug + self.normalize = normalize + self.encoder = encoder + + if not hasattr(self.encoder, "interctc_use_conditioning"): + self.encoder.interctc_use_conditioning = False + if self.encoder.interctc_use_conditioning: + self.encoder.conditioning_layer = torch.nn.Linear( + vocab_size, self.encoder.output_size() + ) + self.interctc_weight = interctc_weight + + # self.error_calculator = None + if ctc_weight == 1.0: + self.decoder = None + else: + self.decoder = decoder + + self.criterion_att = LabelSmoothingLoss( + size=vocab_size, + padding_idx=ignore_id, + smoothing=lsm_weight, + normalize_length=length_normalized_loss, + ) + # + # if report_cer or report_wer: + # self.error_calculator = ErrorCalculator( + # token_list, sym_space, sym_blank, report_cer, report_wer + # ) + # + self.error_calculator = None + if ctc_weight == 0.0: + self.ctc = None + else: + self.ctc = ctc + + self.share_embedding = share_embedding + if self.share_embedding: + self.decoder.embed = None + + self.length_normalized_loss = length_normalized_loss + 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]: + """Encoder + Decoder + Calc loss + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + text: (Batch, Length) + text_lengths: (Batch,) + """ + if len(text_lengths.size()) > 1: + text_lengths = text_lengths[:, 0] + if len(speech_lengths.size()) > 1: + speech_lengths = speech_lengths[:, 0] + + batch_size = speech.shape[0] + + # 1. Encoder + encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) + 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 = None, None, None, None + loss_ctc, cer_ctc = None, None + stats = dict() + + # decoder: CTC branch + if self.ctc_weight != 0.0: + loss_ctc, cer_ctc = self._calc_ctc_loss( + encoder_out, encoder_out_lens, text, text_lengths + ) + + # Collect CTC branch stats + stats["loss_ctc"] = loss_ctc.detach() if loss_ctc is not None else None + stats["cer_ctc"] = cer_ctc + + # Intermediate CTC (optional) + loss_interctc = 0.0 + if self.interctc_weight != 0.0 and intermediate_outs is not None: + for layer_idx, intermediate_out in intermediate_outs: + # we assume intermediate_out has the same length & padding + # as those of encoder_out + loss_ic, cer_ic = self._calc_ctc_loss( + intermediate_out, encoder_out_lens, text, text_lengths + ) + loss_interctc = loss_interctc + loss_ic + + # Collect Intermedaite CTC stats + stats["loss_interctc_layer{}".format(layer_idx)] = ( + loss_ic.detach() if loss_ic is not None else None + ) + stats["cer_interctc_layer{}".format(layer_idx)] = cer_ic + + loss_interctc = loss_interctc / len(intermediate_outs) + + # calculate whole encoder loss + loss_ctc = (1 - self.interctc_weight) * loss_ctc + self.interctc_weight * loss_interctc + + # decoder: Attention decoder branch + loss_att, acc_att, cer_att, wer_att = self._calc_att_loss( + encoder_out, encoder_out_lens, text, text_lengths + ) + + # 3. CTC-Att loss definition + if self.ctc_weight == 0.0: + loss = loss_att + elif self.ctc_weight == 1.0: + loss = loss_ctc + else: + loss = self.ctc_weight * loss_ctc + (1 - self.ctc_weight) * loss_att + + # Collect Attn branch stats + stats["loss_att"] = loss_att.detach() if loss_att is not None else None + stats["acc"] = acc_att + stats["cer"] = cer_att + stats["wer"] = wer_att + + # Collect total loss stats + 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 encode( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Frontend + Encoder. Note that this method is used by asr_inference.py + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + ind: int + """ + with autocast(False): + + # Data augmentation + if self.specaug is not None and self.training: + speech, speech_lengths = self.specaug(speech, speech_lengths) + + # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN + if self.normalize is not None: + speech, speech_lengths = self.normalize(speech, speech_lengths) + + # Forward encoder + # feats: (Batch, Length, Dim) + # -> encoder_out: (Batch, Length2, Dim2) + if self.encoder.interctc_use_conditioning: + encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths, ctc=self.ctc) + else: + encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths) + intermediate_outs = None + if isinstance(encoder_out, tuple): + intermediate_outs = encoder_out[1] + encoder_out = encoder_out[0] + + if intermediate_outs is not None: + return (encoder_out, intermediate_outs), encoder_out_lens + + return encoder_out, encoder_out_lens + + def _calc_att_loss( + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, + ): + 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_ctc_loss( + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, + ): + # Calc CTC loss + loss_ctc = self.ctc(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens) + + # Calc CER using CTC + cer_ctc = None + if not self.training and self.error_calculator is not None: + ys_hat = self.ctc.argmax(encoder_out).data + cer_ctc = self.error_calculator(ys_hat.cpu(), ys_pad.cpu(), is_ctc=True) + return loss_ctc, cer_ctc + + def init_beam_search( + self, + **kwargs, + ): + from funasr.models.transformer.search import BeamSearch + from funasr.models.transformer.scorers.ctc import CTCPrefixScorer + from funasr.models.transformer.scorers.length_bonus import LengthBonus + + # 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=self.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.5), + ctc=kwargs.get("decoding_ctc_weight", 0.5), + lm=kwargs.get("lm_weight", 0.0), + ngram=kwargs.get("ngram_weight", 0.0), + length_bonus=kwargs.get("penalty", 0.0), + ) + beam_search = BeamSearch( + beam_size=kwargs.get("beam_size", 10), + 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, + ): + + if kwargs.get("batch_size", 1) > 1: + raise NotImplementedError("batch decoding is not implemented") + + # init beamsearch + if self.beam_search is None: + logging.info("enable beam_search") + self.init_beam_search(**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"]) + # Encoder + encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) + if isinstance(encoder_out, tuple): + encoder_out = encoder_out[0] + + # c. Passed the encoder result and the beam search + nbest_hyps = self.beam_search( + x=encoder_out[0], + maxlenratio=kwargs.get("maxlenratio", 0.0), + minlenratio=kwargs.get("minlenratio", 0.0), + ) + + nbest_hyps = nbest_hyps[: self.nbest] + + results = [] + b, n, d = encoder_out.size() + for i in range(b): + + for nbest_idx, hyp in enumerate(nbest_hyps): + ibest_writer = None + if kwargs.get("output_dir") is not None: + if not hasattr(self, "writer"): + self.writer = DatadirWriter(kwargs.get("output_dir")) + ibest_writer = self.writer[f"{nbest_idx + 1}best_recog"] + + # remove sos/eos and get results + last_pos = -1 + if isinstance(hyp.yseq, list): + token_int = hyp.yseq[1:last_pos] + else: + token_int = hyp.yseq[1:last_pos].tolist() + + # remove blank symbol id, which is assumed to be 0 + token_int = list( + filter( + lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int + ) + ) + + # Change integer-ids to tokens + token = tokenizer.ids2tokens(token_int) + text = tokenizer.tokens2text(token) + + text_postprocessed, _ = postprocess_utils.sentence_postprocess(token) + result_i = {"key": key[i], "token": token, "text": text_postprocessed} + results.append(result_i) + + if ibest_writer is not None: + ibest_writer["token"][key[i]] = " ".join(token) + ibest_writer["text"][key[i]] = text_postprocessed + + return results, meta_data diff --git a/funasr/models/transformer/positionwise_feed_forward.py b/funasr/models/transformer/positionwise_feed_forward.py new file mode 100644 index 0000000000000000000000000000000000000000..7cfa5f9d51991cb0c8f9dfe2b42221dd5b5b9885 --- /dev/null +++ b/funasr/models/transformer/positionwise_feed_forward.py @@ -0,0 +1,48 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +# Copyright 2019 Shigeki Karita +# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +"""Positionwise feed forward layer definition.""" + +import torch + +from funasr.models.transformer.layer_norm import LayerNorm + + +class PositionwiseFeedForward(torch.nn.Module): + """Positionwise feed forward layer. + + Args: + idim (int): Input dimenstion. + hidden_units (int): The number of hidden units. + dropout_rate (float): Dropout rate. + + """ + + def __init__(self, idim, hidden_units, dropout_rate, activation=torch.nn.ReLU()): + """Construct an PositionwiseFeedForward object.""" + super(PositionwiseFeedForward, self).__init__() + self.w_1 = torch.nn.Linear(idim, hidden_units) + self.w_2 = torch.nn.Linear(hidden_units, idim) + self.dropout = torch.nn.Dropout(dropout_rate) + self.activation = activation + + def forward(self, x): + """Forward function.""" + return self.w_2(self.dropout(self.activation(self.w_1(x)))) + + +class PositionwiseFeedForwardDecoderSANMExport(torch.nn.Module): + def __init__(self, model): + super().__init__() + self.w_1 = model.w_1 + self.w_2 = model.w_2 + self.activation = model.activation + self.norm = model.norm + + def forward(self, x): + x = self.activation(self.w_1(x)) + x = self.w_2(self.norm(x)) + return x diff --git a/funasr/models/transformer/scorers/__init__.py b/funasr/models/transformer/scorers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b7f177368e62a5578b8706300e101f831a3972ac --- /dev/null +++ b/funasr/models/transformer/scorers/__init__.py @@ -0,0 +1 @@ +"""Initialize sub package.""" diff --git a/funasr/models/transformer/scorers/ctc.py b/funasr/models/transformer/scorers/ctc.py new file mode 100644 index 0000000000000000000000000000000000000000..eb19d18910da2586456d85060c6081338f98e4ff --- /dev/null +++ b/funasr/models/transformer/scorers/ctc.py @@ -0,0 +1,155 @@ +"""ScorerInterface implementation for CTC.""" + +import numpy as np +import torch + +from funasr.models.transformer.scorers.ctc_prefix_score import CTCPrefixScore +from funasr.models.transformer.scorers.ctc_prefix_score import CTCPrefixScoreTH +from funasr.models.transformer.scorers.scorer_interface import BatchPartialScorerInterface + +class CTCPrefixScorer(BatchPartialScorerInterface): + """Decoder interface wrapper for CTCPrefixScore.""" + + def __init__(self, ctc: torch.nn.Module, eos: int): + """Initialize class. + + Args: + ctc (torch.nn.Module): The CTC implementation. + For example, :class:`espnet.nets.pytorch_backend.ctc.CTC` + eos (int): The end-of-sequence id. + + """ + self.ctc = ctc + self.eos = eos + self.impl = None + + def init_state(self, x: torch.Tensor): + """Get an initial state for decoding. + + Args: + x (torch.Tensor): The encoded feature tensor + + Returns: initial state + + """ + logp = self.ctc.log_softmax(x.unsqueeze(0)).detach().squeeze(0).cpu().numpy() + # TODO(karita): use CTCPrefixScoreTH + self.impl = CTCPrefixScore(logp, 0, self.eos, np) + return 0, self.impl.initial_state() + + def select_state(self, state, i, new_id=None): + """Select state with relative ids in the main beam search. + + Args: + state: Decoder state for prefix tokens + i (int): Index to select a state in the main beam search + new_id (int): New label id to select a state if necessary + + Returns: + state: pruned state + + """ + if type(state) == tuple: + if len(state) == 2: # for CTCPrefixScore + sc, st = state + return sc[i], st[i] + else: # for CTCPrefixScoreTH (need new_id > 0) + r, log_psi, f_min, f_max, scoring_idmap = state + s = log_psi[i, new_id].expand(log_psi.size(1)) + if scoring_idmap is not None: + return r[:, :, i, scoring_idmap[i, new_id]], s, f_min, f_max + else: + return r[:, :, i, new_id], s, f_min, f_max + return None if state is None else state[i] + + def score_partial(self, y, ids, state, x): + """Score new token. + + Args: + y (torch.Tensor): 1D prefix token + next_tokens (torch.Tensor): torch.int64 next token to score + state: decoder state for prefix tokens + x (torch.Tensor): 2D encoder feature that generates ys + + Returns: + tuple[torch.Tensor, Any]: + Tuple of a score tensor for y that has a shape `(len(next_tokens),)` + and next state for ys + + """ + prev_score, state = state + presub_score, new_st = self.impl(y.cpu(), ids.cpu(), state) + tscore = torch.as_tensor(presub_score - prev_score, device=x.device, dtype=x.dtype) + return tscore, (presub_score, new_st) + + def batch_init_state(self, x: torch.Tensor): + """Get an initial state for decoding. + + Args: + x (torch.Tensor): The encoded feature tensor + + Returns: initial state + + """ + logp = self.ctc.log_softmax(x.unsqueeze(0)) # assuming batch_size = 1 + xlen = torch.tensor([logp.size(1)]) + self.impl = CTCPrefixScoreTH(logp, xlen, 0, self.eos) + return None + + def batch_score_partial(self, y, ids, state, x): + """Score new token. + + Args: + y (torch.Tensor): 1D prefix token + ids (torch.Tensor): torch.int64 next token to score + state: decoder state for prefix tokens + x (torch.Tensor): 2D encoder feature that generates ys + + Returns: + tuple[torch.Tensor, Any]: + Tuple of a score tensor for y that has a shape `(len(next_tokens),)` + and next state for ys + + """ + batch_state = ( + ( + torch.stack([s[0] for s in state], dim=2), + torch.stack([s[1] for s in state]), + state[0][2], + state[0][3], + ) + if state[0] is not None + else None + ) + return self.impl(y, batch_state, ids) + + def extend_prob(self, x: torch.Tensor): + """Extend probs for decoding. + + This extension is for streaming decoding + as in Eq (14) in https://arxiv.org/abs/2006.14941 + + Args: + x (torch.Tensor): The encoded feature tensor + + """ + logp = self.ctc.log_softmax(x.unsqueeze(0)) + self.impl.extend_prob(logp) + + def extend_state(self, state): + """Extend state for decoding. + + This extension is for streaming decoding + as in Eq (14) in https://arxiv.org/abs/2006.14941 + + Args: + state: The states of hyps + + Returns: exteded state + + """ + new_state = [] + for s in state: + new_state.append(self.impl.extend_state(s)) + + return new_state diff --git a/funasr/models/transformer/scorers/ctc_prefix_score.py b/funasr/models/transformer/scorers/ctc_prefix_score.py new file mode 100644 index 0000000000000000000000000000000000000000..ce0703e736b35b98d75b2893d4aac589b99a6a8b --- /dev/null +++ b/funasr/models/transformer/scorers/ctc_prefix_score.py @@ -0,0 +1,337 @@ +#!/usr/bin/env python3 + +# Copyright 2018 Mitsubishi Electric Research Labs (Takaaki Hori) +# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +import torch + +import numpy as np +import six + + +class CTCPrefixScoreTH(object): + """Batch processing of CTCPrefixScore + + which is based on Algorithm 2 in WATANABE et al. + "HYBRID CTC/ATTENTION ARCHITECTURE FOR END-TO-END SPEECH RECOGNITION," + but extended to efficiently compute the label probablities for multiple + hypotheses simultaneously + See also Seki et al. "Vectorized Beam Search for CTC-Attention-Based + Speech Recognition," In INTERSPEECH (pp. 3825-3829), 2019. + """ + + def __init__(self, x, xlens, blank, eos, margin=0): + """Construct CTC prefix scorer + + :param torch.Tensor x: input label posterior sequences (B, T, O) + :param torch.Tensor xlens: input lengths (B,) + :param int blank: blank label id + :param int eos: end-of-sequence id + :param int margin: margin parameter for windowing (0 means no windowing) + """ + # In the comment lines, + # we assume T: input_length, B: batch size, W: beam width, O: output dim. + self.logzero = -10000000000.0 + self.blank = blank + self.eos = eos + self.batch = x.size(0) + self.input_length = x.size(1) + self.odim = x.size(2) + self.dtype = x.dtype + self.device = torch.device("cuda:%d" % x.get_device()) if x.is_cuda else torch.device("cpu") + # Pad the rest of posteriors in the batch + # TODO(takaaki-hori): need a better way without for-loops + for i, l in enumerate(xlens): + if l < self.input_length: + x[i, l:, :] = self.logzero + x[i, l:, blank] = 0 + # Reshape input x + xn = x.transpose(0, 1) # (B, T, O) -> (T, B, O) + xb = xn[:, :, self.blank].unsqueeze(2).expand(-1, -1, self.odim) + self.x = torch.stack([xn, xb]) # (2, T, B, O) + self.end_frames = torch.as_tensor(xlens) - 1 + + # Setup CTC windowing + self.margin = margin + if margin > 0: + self.frame_ids = torch.arange(self.input_length, dtype=self.dtype, device=self.device) + # Base indices for index conversion + self.idx_bh = None + self.idx_b = torch.arange(self.batch, device=self.device) + self.idx_bo = (self.idx_b * self.odim).unsqueeze(1) + + def __call__(self, y, state, scoring_ids=None, att_w=None): + """Compute CTC prefix scores for next labels + + :param list y: prefix label sequences + :param tuple state: previous CTC state + :param torch.Tensor pre_scores: scores for pre-selection of hypotheses (BW, O) + :param torch.Tensor att_w: attention weights to decide CTC window + :return new_state, ctc_local_scores (BW, O) + """ + output_length = len(y[0]) - 1 # ignore sos + last_ids = [yi[-1] for yi in y] # last output label ids + n_bh = len(last_ids) # batch * hyps + n_hyps = n_bh // self.batch # assuming each utterance has the same # of hyps + self.scoring_num = scoring_ids.size(-1) if scoring_ids is not None else 0 + # prepare state info + if state is None: + r_prev = torch.full( + (self.input_length, 2, self.batch, n_hyps), + self.logzero, + dtype=self.dtype, + device=self.device, + ) + r_prev[:, 1] = torch.cumsum(self.x[0, :, :, self.blank], 0).unsqueeze(2) + r_prev = r_prev.view(-1, 2, n_bh) + s_prev = 0.0 + f_min_prev = 0 + f_max_prev = 1 + else: + r_prev, s_prev, f_min_prev, f_max_prev = state + + # select input dimensions for scoring + if self.scoring_num > 0: + scoring_idmap = torch.full((n_bh, self.odim), -1, dtype=torch.long, device=self.device) + snum = self.scoring_num + if self.idx_bh is None or n_bh > len(self.idx_bh): + self.idx_bh = torch.arange(n_bh, device=self.device).view(-1, 1) + scoring_idmap[self.idx_bh[:n_bh], scoring_ids] = torch.arange(snum, device=self.device) + scoring_idx = (scoring_ids + self.idx_bo.repeat(1, n_hyps).view(-1, 1)).view(-1) + x_ = torch.index_select( + self.x.view(2, -1, self.batch * self.odim), 2, scoring_idx + ).view(2, -1, n_bh, snum) + else: + scoring_ids = None + scoring_idmap = None + snum = self.odim + x_ = self.x.unsqueeze(3).repeat(1, 1, 1, n_hyps, 1).view(2, -1, n_bh, snum) + + # new CTC forward probs are prepared as a (T x 2 x BW x S) tensor + # that corresponds to r_t^n(h) and r_t^b(h) in a batch. + r = torch.full( + (self.input_length, 2, n_bh, snum), + self.logzero, + dtype=self.dtype, + device=self.device, + ) + if output_length == 0: + r[0, 0] = x_[0, 0] + + r_sum = torch.logsumexp(r_prev, 1) + log_phi = r_sum.unsqueeze(2).repeat(1, 1, snum) + if scoring_ids is not None: + for idx in range(n_bh): + pos = scoring_idmap[idx, last_ids[idx]] + if pos >= 0: + log_phi[:, idx, pos] = r_prev[:, 1, idx] + else: + for idx in range(n_bh): + log_phi[:, idx, last_ids[idx]] = r_prev[:, 1, idx] + + # decide start and end frames based on attention weights + if att_w is not None and self.margin > 0: + f_arg = torch.matmul(att_w, self.frame_ids) + f_min = max(int(f_arg.min().cpu()), f_min_prev) + f_max = max(int(f_arg.max().cpu()), f_max_prev) + start = min(f_max_prev, max(f_min - self.margin, output_length, 1)) + end = min(f_max + self.margin, self.input_length) + else: + f_min = f_max = 0 + start = max(output_length, 1) + end = self.input_length + + # compute forward probabilities log(r_t^n(h)) and log(r_t^b(h)) + for t in range(start, end): + rp = r[t - 1] + rr = torch.stack([rp[0], log_phi[t - 1], rp[0], rp[1]]).view(2, 2, n_bh, snum) + r[t] = torch.logsumexp(rr, 1) + x_[:, t] + + # compute log prefix probabilities log(psi) + log_phi_x = torch.cat((log_phi[0].unsqueeze(0), log_phi[:-1]), dim=0) + x_[0] + if scoring_ids is not None: + log_psi = torch.full( + (n_bh, self.odim), self.logzero, dtype=self.dtype, device=self.device + ) + log_psi_ = torch.logsumexp( + torch.cat((log_phi_x[start:end], r[start - 1, 0].unsqueeze(0)), dim=0), + dim=0, + ) + for si in range(n_bh): + log_psi[si, scoring_ids[si]] = log_psi_[si] + else: + log_psi = torch.logsumexp( + torch.cat((log_phi_x[start:end], r[start - 1, 0].unsqueeze(0)), dim=0), + dim=0, + ) + + for si in range(n_bh): + log_psi[si, self.eos] = r_sum[self.end_frames[si // n_hyps], si] + + # exclude blank probs + log_psi[:, self.blank] = self.logzero + + return (log_psi - s_prev), (r, log_psi, f_min, f_max, scoring_idmap) + + def index_select_state(self, state, best_ids): + """Select CTC states according to best ids + + :param state : CTC state + :param best_ids : index numbers selected by beam pruning (B, W) + :return selected_state + """ + r, s, f_min, f_max, scoring_idmap = state + # convert ids to BHO space + n_bh = len(s) + n_hyps = n_bh // self.batch + vidx = (best_ids + (self.idx_b * (n_hyps * self.odim)).view(-1, 1)).view(-1) + # select hypothesis scores + s_new = torch.index_select(s.view(-1), 0, vidx) + s_new = s_new.view(-1, 1).repeat(1, self.odim).view(n_bh, self.odim) + # convert ids to BHS space (S: scoring_num) + if scoring_idmap is not None: + snum = self.scoring_num + hyp_idx = (best_ids // self.odim + (self.idx_b * n_hyps).view(-1, 1)).view(-1) + label_ids = torch.fmod(best_ids, self.odim).view(-1) + score_idx = scoring_idmap[hyp_idx, label_ids] + score_idx[score_idx == -1] = 0 + vidx = score_idx + hyp_idx * snum + else: + snum = self.odim + # select forward probabilities + r_new = torch.index_select(r.view(-1, 2, n_bh * snum), 2, vidx).view(-1, 2, n_bh) + return r_new, s_new, f_min, f_max + + def extend_prob(self, x): + """Extend CTC prob. + + :param torch.Tensor x: input label posterior sequences (B, T, O) + """ + + if self.x.shape[1] < x.shape[1]: # self.x (2,T,B,O); x (B,T,O) + # Pad the rest of posteriors in the batch + # TODO(takaaki-hori): need a better way without for-loops + xlens = [x.size(1)] + for i, l in enumerate(xlens): + if l < self.input_length: + x[i, l:, :] = self.logzero + x[i, l:, self.blank] = 0 + tmp_x = self.x + xn = x.transpose(0, 1) # (B, T, O) -> (T, B, O) + xb = xn[:, :, self.blank].unsqueeze(2).expand(-1, -1, self.odim) + self.x = torch.stack([xn, xb]) # (2, T, B, O) + self.x[:, : tmp_x.shape[1], :, :] = tmp_x + self.input_length = x.size(1) + self.end_frames = torch.as_tensor(xlens) - 1 + + def extend_state(self, state): + """Compute CTC prefix state. + + + :param state : CTC state + :return ctc_state + """ + + if state is None: + # nothing to do + return state + else: + r_prev, s_prev, f_min_prev, f_max_prev = state + + r_prev_new = torch.full( + (self.input_length, 2), + self.logzero, + dtype=self.dtype, + device=self.device, + ) + start = max(r_prev.shape[0], 1) + r_prev_new[0:start] = r_prev + for t in six.moves.range(start, self.input_length): + r_prev_new[t, 1] = r_prev_new[t - 1, 1] + self.x[0, t, :, self.blank] + + return (r_prev_new, s_prev, f_min_prev, f_max_prev) + + +class CTCPrefixScore(object): + """Compute CTC label sequence scores + + which is based on Algorithm 2 in WATANABE et al. + "HYBRID CTC/ATTENTION ARCHITECTURE FOR END-TO-END SPEECH RECOGNITION," + but extended to efficiently compute the probablities of multiple labels + simultaneously + """ + + def __init__(self, x, blank, eos, xp): + self.xp = xp + self.logzero = -10000000000.0 + self.blank = blank + self.eos = eos + self.input_length = len(x) + self.x = x + + def initial_state(self): + """Obtain an initial CTC state + + :return: CTC state + """ + # initial CTC state is made of a frame x 2 tensor that corresponds to + # r_t^n() and r_t^b(), where 0 and 1 of axis=1 represent + # superscripts n and b (non-blank and blank), respectively. + r = self.xp.full((self.input_length, 2), self.logzero, dtype=np.float32) + r[0, 1] = self.x[0, self.blank] + for i in six.moves.range(1, self.input_length): + r[i, 1] = r[i - 1, 1] + self.x[i, self.blank] + return r + + def __call__(self, y, cs, r_prev): + """Compute CTC prefix scores for next labels + + :param y : prefix label sequence + :param cs : array of next labels + :param r_prev: previous CTC state + :return ctc_scores, ctc_states + """ + # initialize CTC states + output_length = len(y) - 1 # ignore sos + # new CTC states are prepared as a frame x (n or b) x n_labels tensor + # that corresponds to r_t^n(h) and r_t^b(h). + r = self.xp.ndarray((self.input_length, 2, len(cs)), dtype=np.float32) + xs = self.x[:, cs] + if output_length == 0: + r[0, 0] = xs[0] + r[0, 1] = self.logzero + else: + r[output_length - 1] = self.logzero + + # prepare forward probabilities for the last label + r_sum = self.xp.logaddexp(r_prev[:, 0], r_prev[:, 1]) # log(r_t^n(g) + r_t^b(g)) + last = y[-1] + if output_length > 0 and last in cs: + log_phi = self.xp.ndarray((self.input_length, len(cs)), dtype=np.float32) + for i in six.moves.range(len(cs)): + log_phi[:, i] = r_sum if cs[i] != last else r_prev[:, 1] + else: + log_phi = r_sum + + # compute forward probabilities log(r_t^n(h)), log(r_t^b(h)), + # and log prefix probabilities log(psi) + start = max(output_length, 1) + log_psi = r[start - 1, 0] + for t in six.moves.range(start, self.input_length): + r[t, 0] = self.xp.logaddexp(r[t - 1, 0], log_phi[t - 1]) + xs[t] + r[t, 1] = self.xp.logaddexp(r[t - 1, 0], r[t - 1, 1]) + self.x[t, self.blank] + log_psi = self.xp.logaddexp(log_psi, log_phi[t - 1] + xs[t]) + + # get P(...eos|X) that ends with the prefix itself + eos_pos = self.xp.where(cs == self.eos)[0] + if len(eos_pos) > 0: + log_psi[eos_pos] = r_sum[-1] # log(r_T^n(g) + r_T^b(g)) + + # exclude blank probs + blank_pos = self.xp.where(cs == self.blank)[0] + if len(blank_pos) > 0: + log_psi[blank_pos] = self.logzero + + # return the log prefix probability and CTC states, where the label axis + # of the CTC states is moved to the first axis to slice it easily + return log_psi, self.xp.rollaxis(r, 2) diff --git a/funasr/models/transformer/scorers/length_bonus.py b/funasr/models/transformer/scorers/length_bonus.py new file mode 100644 index 0000000000000000000000000000000000000000..3a7e9e073286c2be480d136409b377e7aced4125 --- /dev/null +++ b/funasr/models/transformer/scorers/length_bonus.py @@ -0,0 +1,60 @@ +"""Length bonus module.""" + +from typing import Any +from typing import List +from typing import Tuple + +import torch + +from funasr.models.transformer.scorers.scorer_interface import BatchScorerInterface + + +class LengthBonus(BatchScorerInterface): + """Length bonus in beam search.""" + + def __init__(self, n_vocab: int): + """Initialize class. + + Args: + n_vocab (int): The number of tokens in vocabulary for beam search + + """ + self.n = n_vocab + + def score(self, y, state, x): + """Score new token. + + Args: + y (torch.Tensor): 1D torch.int64 prefix tokens. + state: Scorer state for prefix tokens + x (torch.Tensor): 2D encoder feature that generates ys. + + Returns: + tuple[torch.Tensor, Any]: Tuple of + torch.float32 scores for next token (n_vocab) + and None + + """ + return torch.tensor([1.0], device=x.device, dtype=x.dtype).expand(self.n), None + + def batch_score( + self, ys: torch.Tensor, states: List[Any], xs: torch.Tensor + ) -> Tuple[torch.Tensor, List[Any]]: + """Score new token batch. + + Args: + ys (torch.Tensor): torch.int64 prefix tokens (n_batch, ylen). + states (List[Any]): Scorer states for prefix tokens. + xs (torch.Tensor): + The encoder feature that generates ys (n_batch, xlen, n_feat). + + Returns: + tuple[torch.Tensor, List[Any]]: Tuple of + batchfied scores for next token with shape of `(n_batch, n_vocab)` + and next state list for ys. + + """ + return ( + torch.tensor([1.0], device=xs.device, dtype=xs.dtype).expand(ys.shape[0], self.n), + None, + ) diff --git a/funasr/models/transformer/scorers/scorer_interface.py b/funasr/models/transformer/scorers/scorer_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..3acdf2f4d1dce506f9f2b9bdf18ec7f4b6341205 --- /dev/null +++ b/funasr/models/transformer/scorers/scorer_interface.py @@ -0,0 +1,186 @@ +"""Scorer interface module.""" + +from typing import Any +from typing import List +from typing import Tuple + +import torch +import warnings + + +class ScorerInterface: + """Scorer interface for beam search. + + The scorer performs scoring of the all tokens in vocabulary. + + Examples: + * Search heuristics + * :class:`espnet.nets.scorers.length_bonus.LengthBonus` + * Decoder networks of the sequence-to-sequence models + * :class:`espnet.nets.pytorch_backend.nets.transformer.decoder.Decoder` + * :class:`espnet.nets.pytorch_backend.nets.rnn.decoders.Decoder` + * Neural language models + * :class:`espnet.nets.pytorch_backend.lm.transformer.TransformerLM` + * :class:`espnet.nets.pytorch_backend.lm.default.DefaultRNNLM` + * :class:`espnet.nets.pytorch_backend.lm.seq_rnn.SequentialRNNLM` + + """ + + def init_state(self, x: torch.Tensor) -> Any: + """Get an initial state for decoding (optional). + + Args: + x (torch.Tensor): The encoded feature tensor + + Returns: initial state + + """ + return None + + def select_state(self, state: Any, i: int, new_id: int = None) -> Any: + """Select state with relative ids in the main beam search. + + Args: + state: Decoder state for prefix tokens + i (int): Index to select a state in the main beam search + new_id (int): New label index to select a state if necessary + + Returns: + state: pruned state + + """ + return None if state is None else state[i] + + def score(self, y: torch.Tensor, state: Any, x: torch.Tensor) -> Tuple[torch.Tensor, Any]: + """Score new token (required). + + Args: + y (torch.Tensor): 1D torch.int64 prefix tokens. + state: Scorer state for prefix tokens + x (torch.Tensor): The encoder feature that generates ys. + + Returns: + tuple[torch.Tensor, Any]: Tuple of + scores for next token that has a shape of `(n_vocab)` + and next state for ys + + """ + raise NotImplementedError + + def final_score(self, state: Any) -> float: + """Score eos (optional). + + Args: + state: Scorer state for prefix tokens + + Returns: + float: final score + + """ + return 0.0 + + +class BatchScorerInterface(ScorerInterface): + """Batch scorer interface.""" + + def batch_init_state(self, x: torch.Tensor) -> Any: + """Get an initial state for decoding (optional). + + Args: + x (torch.Tensor): The encoded feature tensor + + Returns: initial state + + """ + return self.init_state(x) + + def batch_score( + self, ys: torch.Tensor, states: List[Any], xs: torch.Tensor + ) -> Tuple[torch.Tensor, List[Any]]: + """Score new token batch (required). + + Args: + ys (torch.Tensor): torch.int64 prefix tokens (n_batch, ylen). + states (List[Any]): Scorer states for prefix tokens. + xs (torch.Tensor): + The encoder feature that generates ys (n_batch, xlen, n_feat). + + Returns: + tuple[torch.Tensor, List[Any]]: Tuple of + batchfied scores for next token with shape of `(n_batch, n_vocab)` + and next state list for ys. + + """ + warnings.warn( + "{} batch score is implemented through for loop not parallelized".format( + self.__class__.__name__ + ) + ) + scores = list() + outstates = list() + for i, (y, state, x) in enumerate(zip(ys, states, xs)): + score, outstate = self.score(y, state, x) + outstates.append(outstate) + scores.append(score) + scores = torch.cat(scores, 0).view(ys.shape[0], -1) + return scores, outstates + + +class PartialScorerInterface(ScorerInterface): + """Partial scorer interface for beam search. + + The partial scorer performs scoring when non-partial scorer finished scoring, + and receives pre-pruned next tokens to score because it is too heavy to score + all the tokens. + + Examples: + * Prefix search for connectionist-temporal-classification models + * :class:`espnet.nets.scorers.ctc.CTCPrefixScorer` + + """ + + def score_partial( + self, y: torch.Tensor, next_tokens: torch.Tensor, state: Any, x: torch.Tensor + ) -> Tuple[torch.Tensor, Any]: + """Score new token (required). + + Args: + y (torch.Tensor): 1D prefix token + next_tokens (torch.Tensor): torch.int64 next token to score + state: decoder state for prefix tokens + x (torch.Tensor): The encoder feature that generates ys + + Returns: + tuple[torch.Tensor, Any]: + Tuple of a score tensor for y that has a shape `(len(next_tokens),)` + and next state for ys + + """ + raise NotImplementedError + + +class BatchPartialScorerInterface(BatchScorerInterface, PartialScorerInterface): + """Batch partial scorer interface for beam search.""" + + def batch_score_partial( + self, + ys: torch.Tensor, + next_tokens: torch.Tensor, + states: List[Any], + xs: torch.Tensor, + ) -> Tuple[torch.Tensor, Any]: + """Score new token (required). + + Args: + ys (torch.Tensor): torch.int64 prefix tokens (n_batch, ylen). + next_tokens (torch.Tensor): torch.int64 tokens to score (n_batch, n_token). + states (List[Any]): Scorer states for prefix tokens. + xs (torch.Tensor): + The encoder feature that generates ys (n_batch, xlen, n_feat). + + Returns: + tuple[torch.Tensor, Any]: + Tuple of a score tensor for ys that has a shape `(n_batch, n_vocab)` + and next states for ys + """ + raise NotImplementedError diff --git a/funasr/models/transformer/search.py b/funasr/models/transformer/search.py new file mode 100644 index 0000000000000000000000000000000000000000..c1d90b64e4b322ba1fb093f0db763a63328819a8 --- /dev/null +++ b/funasr/models/transformer/search.py @@ -0,0 +1,449 @@ +from itertools import chain +import logging +from typing import Any +from typing import Dict +from typing import List +from typing import NamedTuple +from typing import Tuple +from typing import Union + +import torch + +from funasr.metrics.common import end_detect +from funasr.models.transformer.scorers.scorer_interface import PartialScorerInterface +from funasr.models.transformer.scorers.scorer_interface import ScorerInterface + + +class Hypothesis(NamedTuple): + """Hypothesis data type.""" + + yseq: torch.Tensor + score: Union[float, torch.Tensor] = 0 + scores: Dict[str, Union[float, torch.Tensor]] = dict() + states: Dict[str, Any] = dict() + + def asdict(self) -> dict: + """Convert data to JSON-friendly dict.""" + return self._replace( + yseq=self.yseq.tolist(), + score=float(self.score), + scores={k: float(v) for k, v in self.scores.items()}, + )._asdict() + + +class BeamSearch(torch.nn.Module): + """Beam search implementation.""" + + def __init__( + self, + scorers: Dict[str, ScorerInterface], + weights: Dict[str, float], + beam_size: int, + vocab_size: int, + sos: int, + eos: int, + token_list: List[str] = None, + pre_beam_ratio: float = 1.5, + pre_beam_score_key: str = None, + ): + """Initialize beam search. + + Args: + scorers (dict[str, ScorerInterface]): Dict of decoder modules + e.g., Decoder, CTCPrefixScorer, LM + The scorer will be ignored if it is `None` + weights (dict[str, float]): Dict of weights for each scorers + The scorer will be ignored if its weight is 0 + beam_size (int): The number of hypotheses kept during search + vocab_size (int): The number of vocabulary + sos (int): Start of sequence id + eos (int): End of sequence id + token_list (list[str]): List of tokens for debug log + pre_beam_score_key (str): key of scores to perform pre-beam search + pre_beam_ratio (float): beam size in the pre-beam search + will be `int(pre_beam_ratio * beam_size)` + + """ + super().__init__() + # set scorers + self.weights = weights + self.scorers = dict() + self.full_scorers = dict() + self.part_scorers = dict() + # this module dict is required for recursive cast + # `self.to(device, dtype)` in `recog.py` + self.nn_dict = torch.nn.ModuleDict() + for k, v in scorers.items(): + w = weights.get(k, 0) + if w == 0 or v is None: + continue + assert isinstance( + v, ScorerInterface + ), f"{k} ({type(v)}) does not implement ScorerInterface" + self.scorers[k] = v + if isinstance(v, PartialScorerInterface): + self.part_scorers[k] = v + else: + self.full_scorers[k] = v + if isinstance(v, torch.nn.Module): + self.nn_dict[k] = v + + # set configurations + self.sos = sos + self.eos = eos + self.token_list = token_list + self.pre_beam_size = int(pre_beam_ratio * beam_size) + self.beam_size = beam_size + self.n_vocab = vocab_size + if ( + pre_beam_score_key is not None + and pre_beam_score_key != "full" + and pre_beam_score_key not in self.full_scorers + ): + raise KeyError(f"{pre_beam_score_key} is not found in {self.full_scorers}") + self.pre_beam_score_key = pre_beam_score_key + self.do_pre_beam = ( + self.pre_beam_score_key is not None + and self.pre_beam_size < self.n_vocab + and len(self.part_scorers) > 0 + ) + + def init_hyp(self, x: torch.Tensor) -> List[Hypothesis]: + """Get an initial hypothesis data. + + Args: + x (torch.Tensor): The encoder output feature + + Returns: + Hypothesis: The initial hypothesis. + + """ + init_states = dict() + init_scores = dict() + for k, d in self.scorers.items(): + init_states[k] = d.init_state(x) + init_scores[k] = 0.0 + return [ + Hypothesis( + score=0.0, + scores=init_scores, + states=init_states, + yseq=torch.tensor([self.sos], device=x.device), + ) + ] + + @staticmethod + def append_token(xs: torch.Tensor, x: int) -> torch.Tensor: + """Append new token to prefix tokens. + + Args: + xs (torch.Tensor): The prefix token + x (int): The new token to append + + Returns: + torch.Tensor: New tensor contains: xs + [x] with xs.dtype and xs.device + + """ + x = torch.tensor([x], dtype=xs.dtype, device=xs.device) + return torch.cat((xs, x)) + + def score_full( + self, hyp: Hypothesis, x: torch.Tensor + ) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: + """Score new hypothesis by `self.full_scorers`. + + Args: + hyp (Hypothesis): Hypothesis with prefix tokens to score + x (torch.Tensor): Corresponding input feature + + Returns: + Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: Tuple of + score dict of `hyp` that has string keys of `self.full_scorers` + and tensor score values of shape: `(self.n_vocab,)`, + and state dict that has string keys + and state values of `self.full_scorers` + + """ + scores = dict() + states = dict() + for k, d in self.full_scorers.items(): + scores[k], states[k] = d.score(hyp.yseq, hyp.states[k], x) + return scores, states + + def score_partial( + self, hyp: Hypothesis, ids: torch.Tensor, x: torch.Tensor + ) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: + """Score new hypothesis by `self.part_scorers`. + + Args: + hyp (Hypothesis): Hypothesis with prefix tokens to score + ids (torch.Tensor): 1D tensor of new partial tokens to score + x (torch.Tensor): Corresponding input feature + + Returns: + Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: Tuple of + score dict of `hyp` that has string keys of `self.part_scorers` + and tensor score values of shape: `(len(ids),)`, + and state dict that has string keys + and state values of `self.part_scorers` + + """ + scores = dict() + states = dict() + for k, d in self.part_scorers.items(): + scores[k], states[k] = d.score_partial(hyp.yseq, ids, hyp.states[k], x) + return scores, states + + def beam( + self, weighted_scores: torch.Tensor, ids: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute topk full token ids and partial token ids. + + Args: + weighted_scores (torch.Tensor): The weighted sum scores for each tokens. + Its shape is `(self.n_vocab,)`. + ids (torch.Tensor): The partial token ids to compute topk + + Returns: + Tuple[torch.Tensor, torch.Tensor]: + The topk full token ids and partial token ids. + Their shapes are `(self.beam_size,)` + + """ + # no pre beam performed + if weighted_scores.size(0) == ids.size(0): + top_ids = weighted_scores.topk(self.beam_size)[1] + return top_ids, top_ids + + # mask pruned in pre-beam not to select in topk + tmp = weighted_scores[ids] + weighted_scores[:] = -float("inf") + weighted_scores[ids] = tmp + top_ids = weighted_scores.topk(self.beam_size)[1] + local_ids = weighted_scores[ids].topk(self.beam_size)[1] + return top_ids, local_ids + + @staticmethod + def merge_scores( + prev_scores: Dict[str, float], + next_full_scores: Dict[str, torch.Tensor], + full_idx: int, + next_part_scores: Dict[str, torch.Tensor], + part_idx: int, + ) -> Dict[str, torch.Tensor]: + """Merge scores for new hypothesis. + + Args: + prev_scores (Dict[str, float]): + The previous hypothesis scores by `self.scorers` + next_full_scores (Dict[str, torch.Tensor]): scores by `self.full_scorers` + full_idx (int): The next token id for `next_full_scores` + next_part_scores (Dict[str, torch.Tensor]): + scores of partial tokens by `self.part_scorers` + part_idx (int): The new token id for `next_part_scores` + + Returns: + Dict[str, torch.Tensor]: The new score dict. + Its keys are names of `self.full_scorers` and `self.part_scorers`. + Its values are scalar tensors by the scorers. + + """ + new_scores = dict() + for k, v in next_full_scores.items(): + new_scores[k] = prev_scores[k] + v[full_idx] + for k, v in next_part_scores.items(): + new_scores[k] = prev_scores[k] + v[part_idx] + return new_scores + + def merge_states(self, states: Any, part_states: Any, part_idx: int) -> Any: + """Merge states for new hypothesis. + + Args: + states: states of `self.full_scorers` + part_states: states of `self.part_scorers` + part_idx (int): The new token id for `part_scores` + + Returns: + Dict[str, torch.Tensor]: The new score dict. + Its keys are names of `self.full_scorers` and `self.part_scorers`. + Its values are states of the scorers. + + """ + new_states = dict() + for k, v in states.items(): + new_states[k] = v + for k, d in self.part_scorers.items(): + new_states[k] = d.select_state(part_states[k], part_idx) + return new_states + + def search(self, running_hyps: List[Hypothesis], x: torch.Tensor) -> List[Hypothesis]: + """Search new tokens for running hypotheses and encoded speech x. + + Args: + running_hyps (List[Hypothesis]): Running hypotheses on beam + x (torch.Tensor): Encoded speech feature (T, D) + + Returns: + List[Hypotheses]: Best sorted hypotheses + + """ + best_hyps = [] + part_ids = torch.arange(self.n_vocab, device=x.device) # no pre-beam + for hyp in running_hyps: + # scoring + weighted_scores = torch.zeros(self.n_vocab, dtype=x.dtype, device=x.device) + scores, states = self.score_full(hyp, x) + for k in self.full_scorers: + weighted_scores += self.weights[k] * scores[k] + # partial scoring + if self.do_pre_beam: + pre_beam_scores = ( + weighted_scores + if self.pre_beam_score_key == "full" + else scores[self.pre_beam_score_key] + ) + part_ids = torch.topk(pre_beam_scores, self.pre_beam_size)[1] + part_scores, part_states = self.score_partial(hyp, part_ids, x) + for k in self.part_scorers: + weighted_scores[part_ids] += self.weights[k] * part_scores[k] + # add previous hyp score + weighted_scores += hyp.score + + # update hyps + for j, part_j in zip(*self.beam(weighted_scores, part_ids)): + # will be (2 x beam at most) + best_hyps.append( + Hypothesis( + score=weighted_scores[j], + yseq=self.append_token(hyp.yseq, j), + scores=self.merge_scores(hyp.scores, scores, j, part_scores, part_j), + states=self.merge_states(states, part_states, part_j), + ) + ) + + # sort and prune 2 x beam -> beam + best_hyps = sorted(best_hyps, key=lambda x: x.score, reverse=True)[ + : min(len(best_hyps), self.beam_size) + ] + return best_hyps + + def forward( + self, x: torch.Tensor, maxlenratio: float = 0.0, minlenratio: float = 0.0 + ) -> List[Hypothesis]: + """Perform beam search. + + Args: + x (torch.Tensor): Encoded speech feature (T, D) + maxlenratio (float): Input length ratio to obtain max output length. + If maxlenratio=0.0 (default), it uses a end-detect function + to automatically find maximum hypothesis lengths + If maxlenratio<0.0, its absolute value is interpreted + as a constant max output length. + minlenratio (float): Input length ratio to obtain min output length. + + Returns: + list[Hypothesis]: N-best decoding results + + """ + # set length bounds + if maxlenratio == 0: + maxlen = x.shape[0] + elif maxlenratio < 0: + maxlen = -1 * int(maxlenratio) + else: + maxlen = max(1, int(maxlenratio * x.size(0))) + minlen = int(minlenratio * x.size(0)) + logging.info("decoder input length: " + str(x.shape[0])) + logging.info("max output length: " + str(maxlen)) + logging.info("min output length: " + str(minlen)) + + # main loop of prefix search + running_hyps = self.init_hyp(x) + ended_hyps = [] + for i in range(maxlen): + logging.debug("position " + str(i)) + best = self.search(running_hyps, x) + # post process of one iteration + running_hyps = self.post_process(i, maxlen, maxlenratio, best, ended_hyps) + # end detection + if maxlenratio == 0.0 and end_detect([h.asdict() for h in ended_hyps], i): + logging.info(f"end detected at {i}") + break + if len(running_hyps) == 0: + logging.info("no hypothesis. Finish decoding.") + break + else: + logging.debug(f"remained hypotheses: {len(running_hyps)}") + + nbest_hyps = sorted(ended_hyps, key=lambda x: x.score, reverse=True) + # check the number of hypotheses reaching to eos + if len(nbest_hyps) == 0: + logging.warning( + "there is no N-best results, perform recognition " "again with smaller minlenratio." + ) + return ( + [] + if minlenratio < 0.1 + else self.forward(x, maxlenratio, max(0.0, minlenratio - 0.1)) + ) + + # report the best result + best = nbest_hyps[0] + for k, v in best.scores.items(): + logging.info(f"{v:6.2f} * {self.weights[k]:3} = {v * self.weights[k]:6.2f} for {k}") + logging.info(f"total log probability: {best.score:.2f}") + logging.info(f"normalized log probability: {best.score / len(best.yseq):.2f}") + logging.info(f"total number of ended hypotheses: {len(nbest_hyps)}") + if self.token_list is not None: + logging.info( + "best hypo: " + "".join([self.token_list[x] for x in best.yseq[1:-1]]) + "\n" + ) + return nbest_hyps + + def post_process( + self, + i: int, + maxlen: int, + maxlenratio: float, + running_hyps: List[Hypothesis], + ended_hyps: List[Hypothesis], + ) -> List[Hypothesis]: + """Perform post-processing of beam search iterations. + + Args: + i (int): The length of hypothesis tokens. + maxlen (int): The maximum length of tokens in beam search. + maxlenratio (int): The maximum length ratio in beam search. + running_hyps (List[Hypothesis]): The running hypotheses in beam search. + ended_hyps (List[Hypothesis]): The ended hypotheses in beam search. + + Returns: + List[Hypothesis]: The new running hypotheses. + + """ + logging.debug(f"the number of running hypotheses: {len(running_hyps)}") + if self.token_list is not None: + logging.debug( + "best hypo: " + "".join([self.token_list[x] for x in running_hyps[0].yseq[1:]]) + ) + # add eos in the final loop to avoid that there are no ended hyps + if i == maxlen - 1: + logging.info("adding in the last position in the loop") + running_hyps = [ + h._replace(yseq=self.append_token(h.yseq, self.eos)) for h in running_hyps + ] + + # add ended hypotheses to a final list, and removed them from current hypotheses + # (this will be a problem, number of hyps < beam) + remained_hyps = [] + for hyp in running_hyps: + if hyp.yseq[-1] == self.eos: + # e.g., Word LM needs to add final score + for k, d in chain(self.full_scorers.items(), self.part_scorers.items()): + s = d.final_score(hyp.states[k]) + hyp.scores[k] += s + hyp = hyp._replace(score=hyp.score + self.weights[k] * s) + ended_hyps.append(hyp) + else: + remained_hyps.append(hyp) + return remained_hyps diff --git a/funasr/models/transformer/template.yaml b/funasr/models/transformer/template.yaml new file mode 100644 index 0000000000000000000000000000000000000000..eb597f9ea4e9ecec9bb7d004d92ad172c8b57ee3 --- /dev/null +++ b/funasr/models/transformer/template.yaml @@ -0,0 +1,110 @@ +# This is an example that demonstrates how to configure a model file. +# You can modify the configuration according to your own requirements. + +# to print the register_table: +# from funasr.register import tables +# tables.print() + +# network architecture +model: Transformer +model_conf: + ctc_weight: 0.3 + lsm_weight: 0.1 # label smoothing option + length_normalized_loss: false + +# encoder +encoder: TransformerEncoder +encoder_conf: + output_size: 256 # dimension of attention + attention_heads: 4 + linear_units: 2048 # the number of units of position-wise feed forward + num_blocks: 12 # the number of encoder blocks + dropout_rate: 0.1 + positional_dropout_rate: 0.1 + attention_dropout_rate: 0.0 + input_layer: conv2d # encoder architecture type + normalize_before: true + +# decoder +decoder: TransformerDecoder +decoder_conf: + attention_heads: 4 + linear_units: 2048 + num_blocks: 6 + dropout_rate: 0.1 + positional_dropout_rate: 0.1 + self_attention_dropout_rate: 0.0 + src_attention_dropout_rate: 0.0 + + +# frontend related +frontend: WavFrontend +frontend_conf: + fs: 16000 + window: hamming + n_mels: 80 + frame_length: 25 + frame_shift: 10 + lfr_m: 1 + lfr_n: 1 + +specaug: SpecAug +specaug_conf: + apply_time_warp: true + time_warp_window: 5 + time_warp_mode: bicubic + apply_freq_mask: true + freq_mask_width_range: + - 0 + - 30 + num_freq_mask: 2 + apply_time_mask: true + time_mask_width_range: + - 0 + - 40 + num_time_mask: 2 + +train_conf: + accum_grad: 1 + grad_clip: 5 + max_epoch: 150 + val_scheduler_criterion: + - valid + - acc + best_model_criterion: + - - valid + - acc + - max + keep_nbest_models: 10 + log_interval: 50 + +optim: adam +optim_conf: + lr: 0.002 +scheduler: warmuplr +scheduler_conf: + warmup_steps: 30000 + +dataset: AudioDataset +dataset_conf: + index_ds: IndexDSJsonl + batch_sampler: BatchSampler + batch_type: example # example or length + batch_size: 1 # if batch_type is example, batch_size is the numbers of samples; if length, batch_size is source_token_len+target_token_len; + max_token_length: 2048 # filter samples if source_token_len+target_token_len > max_token_length, + buffer_size: 500 + shuffle: True + num_workers: 0 + +tokenizer: CharTokenizer +tokenizer_conf: + unk_symbol: + split_with_space: true + + +ctc_conf: + dropout_rate: 0.0 + ctc_type: builtin + reduce: true + ignore_nan_grad: true +normalize: null diff --git a/funasr/models/transformer/utils/__init__.py b/funasr/models/transformer/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/models/transformer/utils/add_sos_eos.py b/funasr/models/transformer/utils/add_sos_eos.py new file mode 100644 index 0000000000000000000000000000000000000000..beead1bc36523741b97f33d3c25a9f9824f6e4b3 --- /dev/null +++ b/funasr/models/transformer/utils/add_sos_eos.py @@ -0,0 +1,31 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +# Copyright 2019 Shigeki Karita +# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +"""Unility functions for Transformer.""" + +import torch +from funasr.models.transformer.utils.nets_utils import pad_list + + +def add_sos_eos(ys_pad, sos, eos, ignore_id): + """Add and labels. + + :param torch.Tensor ys_pad: batch of padded target sequences (B, Lmax) + :param int sos: index of + :param int eos: index of + :param int ignore_id: index of padding + :return: padded tensor (B, Lmax) + :rtype: torch.Tensor + :return: padded tensor (B, Lmax) + :rtype: torch.Tensor + """ + + _sos = ys_pad.new([sos]) + _eos = ys_pad.new([eos]) + ys = [y[y != ignore_id] for y in ys_pad] # parse padded ys + ys_in = [torch.cat([_sos, y], dim=0) for y in ys] + ys_out = [torch.cat([y, _eos], dim=0) for y in ys] + return pad_list(ys_in, eos), pad_list(ys_out, ignore_id) diff --git a/funasr/models/transformer/utils/dynamic_conv.py b/funasr/models/transformer/utils/dynamic_conv.py new file mode 100644 index 0000000000000000000000000000000000000000..f6850384b215c8960f4a025028693a7385e72f11 --- /dev/null +++ b/funasr/models/transformer/utils/dynamic_conv.py @@ -0,0 +1,125 @@ +"""Dynamic Convolution module.""" + +import numpy +import torch +from torch import nn +import torch.nn.functional as F + + +MIN_VALUE = float(numpy.finfo(numpy.float32).min) + + +class DynamicConvolution(nn.Module): + """Dynamic Convolution layer. + + This implementation is based on + https://github.com/pytorch/fairseq/tree/master/fairseq + + Args: + wshare (int): the number of kernel of convolution + n_feat (int): the number of features + dropout_rate (float): dropout_rate + kernel_size (int): kernel size (length) + use_kernel_mask (bool): Use causal mask or not for convolution kernel + use_bias (bool): Use bias term or not. + + """ + + def __init__( + self, + wshare, + n_feat, + dropout_rate, + kernel_size, + use_kernel_mask=False, + use_bias=False, + ): + """Construct Dynamic Convolution layer.""" + super(DynamicConvolution, self).__init__() + + assert n_feat % wshare == 0 + self.wshare = wshare + self.use_kernel_mask = use_kernel_mask + self.dropout_rate = dropout_rate + self.kernel_size = kernel_size + self.attn = None + + # linear -> GLU -- -> lightconv -> linear + # \ / + # Linear + self.linear1 = nn.Linear(n_feat, n_feat * 2) + self.linear2 = nn.Linear(n_feat, n_feat) + self.linear_weight = nn.Linear(n_feat, self.wshare * 1 * kernel_size) + nn.init.xavier_uniform(self.linear_weight.weight) + self.act = nn.GLU() + + # dynamic conv related + self.use_bias = use_bias + if self.use_bias: + self.bias = nn.Parameter(torch.Tensor(n_feat)) + + def forward(self, query, key, value, mask): + """Forward of 'Dynamic Convolution'. + + This function takes query, key and value but uses only quert. + This is just for compatibility with self-attention layer (attention.py) + + Args: + query (torch.Tensor): (batch, time1, d_model) input tensor + key (torch.Tensor): (batch, time2, d_model) NOT USED + value (torch.Tensor): (batch, time2, d_model) NOT USED + mask (torch.Tensor): (batch, time1, time2) mask + + Return: + x (torch.Tensor): (batch, time1, d_model) output + + """ + # linear -> GLU -- -> lightconv -> linear + # \ / + # Linear + x = query + B, T, C = x.size() + H = self.wshare + k = self.kernel_size + + # first liner layer + x = self.linear1(x) + + # GLU activation + x = self.act(x) + + # get kernel of convolution + weight = self.linear_weight(x) # B x T x kH + weight = F.dropout(weight, self.dropout_rate, training=self.training) + weight = weight.view(B, T, H, k).transpose(1, 2).contiguous() # B x H x T x k + weight_new = torch.zeros(B * H * T * (T + k - 1), dtype=weight.dtype) + weight_new = weight_new.view(B, H, T, T + k - 1).fill_(float("-inf")) + weight_new = weight_new.to(x.device) # B x H x T x T+k-1 + weight_new.as_strided((B, H, T, k), ((T + k - 1) * T * H, (T + k - 1) * T, T + k, 1)).copy_( + weight + ) + weight_new = weight_new.narrow(-1, int((k - 1) / 2), T) # B x H x T x T(k) + if self.use_kernel_mask: + kernel_mask = torch.tril(torch.ones(T, T, device=x.device)).unsqueeze(0) + weight_new = weight_new.masked_fill(kernel_mask == 0.0, float("-inf")) + weight_new = F.softmax(weight_new, dim=-1) + self.attn = weight_new + weight_new = weight_new.view(B * H, T, T) + + # convolution + x = x.transpose(1, 2).contiguous() # B x C x T + x = x.view(B * H, int(C / H), T).transpose(1, 2) + x = torch.bmm(weight_new, x) # BH x T x C/H + x = x.transpose(1, 2).contiguous().view(B, C, T) + + if self.use_bias: + x = x + self.bias.view(1, -1, 1) + x = x.transpose(1, 2) # B x T x C + + if mask is not None and not self.use_kernel_mask: + mask = mask.transpose(-1, -2) + x = x.masked_fill(mask == 0, 0.0) + + # second linear layer + x = self.linear2(x) + return x diff --git a/funasr/models/transformer/utils/dynamic_conv2d.py b/funasr/models/transformer/utils/dynamic_conv2d.py new file mode 100644 index 0000000000000000000000000000000000000000..29d4dd8b9f3a0cdcb67e9ee31345b2369867a0fa --- /dev/null +++ b/funasr/models/transformer/utils/dynamic_conv2d.py @@ -0,0 +1,136 @@ +"""Dynamic 2-Dimensional Convolution module.""" + +import numpy +import torch +from torch import nn +import torch.nn.functional as F + + +MIN_VALUE = float(numpy.finfo(numpy.float32).min) + + +class DynamicConvolution2D(nn.Module): + """Dynamic 2-Dimensional Convolution layer. + + This implementation is based on + https://github.com/pytorch/fairseq/tree/master/fairseq + + Args: + wshare (int): the number of kernel of convolution + n_feat (int): the number of features + dropout_rate (float): dropout_rate + kernel_size (int): kernel size (length) + use_kernel_mask (bool): Use causal mask or not for convolution kernel + use_bias (bool): Use bias term or not. + + """ + + def __init__( + self, + wshare, + n_feat, + dropout_rate, + kernel_size, + use_kernel_mask=False, + use_bias=False, + ): + """Construct Dynamic 2-Dimensional Convolution layer.""" + super(DynamicConvolution2D, self).__init__() + + assert n_feat % wshare == 0 + self.wshare = wshare + self.use_kernel_mask = use_kernel_mask + self.dropout_rate = dropout_rate + self.kernel_size = kernel_size + self.padding_size = int(kernel_size / 2) + self.attn_t = None + self.attn_f = None + + # linear -> GLU -- -> lightconv -> linear + # \ / + # Linear + self.linear1 = nn.Linear(n_feat, n_feat * 2) + self.linear2 = nn.Linear(n_feat * 2, n_feat) + self.linear_weight = nn.Linear(n_feat, self.wshare * 1 * kernel_size) + nn.init.xavier_uniform(self.linear_weight.weight) + self.linear_weight_f = nn.Linear(n_feat, kernel_size) + nn.init.xavier_uniform(self.linear_weight_f.weight) + self.act = nn.GLU() + + # dynamic conv related + self.use_bias = use_bias + if self.use_bias: + self.bias = nn.Parameter(torch.Tensor(n_feat)) + + def forward(self, query, key, value, mask): + """Forward of 'Dynamic 2-Dimensional Convolution'. + + This function takes query, key and value but uses only query. + This is just for compatibility with self-attention layer (attention.py) + + Args: + query (torch.Tensor): (batch, time1, d_model) input tensor + key (torch.Tensor): (batch, time2, d_model) NOT USED + value (torch.Tensor): (batch, time2, d_model) NOT USED + mask (torch.Tensor): (batch, time1, time2) mask + + Return: + x (torch.Tensor): (batch, time1, d_model) output + + """ + # linear -> GLU -- -> lightconv -> linear + # \ / + # Linear + x = query + B, T, C = x.size() + H = self.wshare + k = self.kernel_size + + # first liner layer + x = self.linear1(x) + + # GLU activation + x = self.act(x) + + # convolution of frequency axis + weight_f = self.linear_weight_f(x).view(B * T, 1, k) # B x T x k + self.attn_f = weight_f.view(B, T, k).unsqueeze(1) + xf = F.conv1d(x.view(1, B * T, C), weight_f, padding=self.padding_size, groups=B * T) + xf = xf.view(B, T, C) + + # get kernel of convolution + weight = self.linear_weight(x) # B x T x kH + weight = F.dropout(weight, self.dropout_rate, training=self.training) + weight = weight.view(B, T, H, k).transpose(1, 2).contiguous() # B x H x T x k + weight_new = torch.zeros(B * H * T * (T + k - 1), dtype=weight.dtype) + weight_new = weight_new.view(B, H, T, T + k - 1).fill_(float("-inf")) + weight_new = weight_new.to(x.device) # B x H x T x T+k-1 + weight_new.as_strided((B, H, T, k), ((T + k - 1) * T * H, (T + k - 1) * T, T + k, 1)).copy_( + weight + ) + weight_new = weight_new.narrow(-1, int((k - 1) / 2), T) # B x H x T x T(k) + if self.use_kernel_mask: + kernel_mask = torch.tril(torch.ones(T, T, device=x.device)).unsqueeze(0) + weight_new = weight_new.masked_fill(kernel_mask == 0.0, float("-inf")) + weight_new = F.softmax(weight_new, dim=-1) + self.attn_t = weight_new + weight_new = weight_new.view(B * H, T, T) + + # convolution + x = x.transpose(1, 2).contiguous() # B x C x T + x = x.view(B * H, int(C / H), T).transpose(1, 2) + x = torch.bmm(weight_new, x) + x = x.transpose(1, 2).contiguous().view(B, C, T) + + if self.use_bias: + x = x + self.bias.view(1, -1, 1) + x = x.transpose(1, 2) # B x T x C + x = torch.cat((x, xf), -1) # B x T x Cx2 + + if mask is not None and not self.use_kernel_mask: + mask = mask.transpose(-1, -2) + x = x.masked_fill(mask == 0, 0.0) + + # second linear layer + x = self.linear2(x) + return x diff --git a/funasr/models/transformer/utils/lightconv.py b/funasr/models/transformer/utils/lightconv.py new file mode 100644 index 0000000000000000000000000000000000000000..b6d131b5f8fe46ffcaad80c079175b427170fe3b --- /dev/null +++ b/funasr/models/transformer/utils/lightconv.py @@ -0,0 +1,108 @@ +"""Lightweight Convolution Module.""" + +import numpy +import torch +from torch import nn +import torch.nn.functional as F + + +MIN_VALUE = float(numpy.finfo(numpy.float32).min) + + +class LightweightConvolution(nn.Module): + """Lightweight Convolution layer. + + This implementation is based on + https://github.com/pytorch/fairseq/tree/master/fairseq + + Args: + wshare (int): the number of kernel of convolution + n_feat (int): the number of features + dropout_rate (float): dropout_rate + kernel_size (int): kernel size (length) + use_kernel_mask (bool): Use causal mask or not for convolution kernel + use_bias (bool): Use bias term or not. + + """ + + def __init__( + self, + wshare, + n_feat, + dropout_rate, + kernel_size, + use_kernel_mask=False, + use_bias=False, + ): + """Construct Lightweight Convolution layer.""" + super(LightweightConvolution, self).__init__() + + assert n_feat % wshare == 0 + self.wshare = wshare + self.use_kernel_mask = use_kernel_mask + self.dropout_rate = dropout_rate + self.kernel_size = kernel_size + self.padding_size = int(kernel_size / 2) + + # linear -> GLU -> lightconv -> linear + self.linear1 = nn.Linear(n_feat, n_feat * 2) + self.linear2 = nn.Linear(n_feat, n_feat) + self.act = nn.GLU() + + # lightconv related + self.weight = nn.Parameter(torch.Tensor(self.wshare, 1, kernel_size).uniform_(0, 1)) + self.use_bias = use_bias + if self.use_bias: + self.bias = nn.Parameter(torch.Tensor(n_feat)) + + # mask of kernel + kernel_mask0 = torch.zeros(self.wshare, int(kernel_size / 2)) + kernel_mask1 = torch.ones(self.wshare, int(kernel_size / 2 + 1)) + self.kernel_mask = torch.cat((kernel_mask1, kernel_mask0), dim=-1).unsqueeze(1) + + def forward(self, query, key, value, mask): + """Forward of 'Lightweight Convolution'. + + This function takes query, key and value but uses only query. + This is just for compatibility with self-attention layer (attention.py) + + Args: + query (torch.Tensor): (batch, time1, d_model) input tensor + key (torch.Tensor): (batch, time2, d_model) NOT USED + value (torch.Tensor): (batch, time2, d_model) NOT USED + mask (torch.Tensor): (batch, time1, time2) mask + + Return: + x (torch.Tensor): (batch, time1, d_model) output + + """ + # linear -> GLU -> lightconv -> linear + x = query + B, T, C = x.size() + H = self.wshare + + # first liner layer + x = self.linear1(x) + + # GLU activation + x = self.act(x) + + # lightconv + x = x.transpose(1, 2).contiguous().view(-1, H, T) # B x C x T + weight = F.dropout(self.weight, self.dropout_rate, training=self.training) + if self.use_kernel_mask: + self.kernel_mask = self.kernel_mask.to(x.device) + weight = weight.masked_fill(self.kernel_mask == 0.0, float("-inf")) + weight = F.softmax(weight, dim=-1) + x = F.conv1d(x, weight, padding=self.padding_size, groups=self.wshare).view(B, C, T) + if self.use_bias: + x = x + self.bias.view(1, -1, 1) + x = x.transpose(1, 2) # B x T x C + + if mask is not None and not self.use_kernel_mask: + mask = mask.transpose(-1, -2) + x = x.masked_fill(mask == 0, 0.0) + + # second linear layer + x = self.linear2(x) + return x diff --git a/funasr/models/transformer/utils/lightconv2d.py b/funasr/models/transformer/utils/lightconv2d.py new file mode 100644 index 0000000000000000000000000000000000000000..60e5af53e6afc865f07cc46830b0bf81646a9f82 --- /dev/null +++ b/funasr/models/transformer/utils/lightconv2d.py @@ -0,0 +1,120 @@ +"""Lightweight 2-Dimensional Convolution module.""" + +import numpy +import torch +from torch import nn +import torch.nn.functional as F + + +MIN_VALUE = float(numpy.finfo(numpy.float32).min) + + +class LightweightConvolution2D(nn.Module): + """Lightweight 2-Dimensional Convolution layer. + + This implementation is based on + https://github.com/pytorch/fairseq/tree/master/fairseq + + Args: + wshare (int): the number of kernel of convolution + n_feat (int): the number of features + dropout_rate (float): dropout_rate + kernel_size (int): kernel size (length) + use_kernel_mask (bool): Use causal mask or not for convolution kernel + use_bias (bool): Use bias term or not. + + """ + + def __init__( + self, + wshare, + n_feat, + dropout_rate, + kernel_size, + use_kernel_mask=False, + use_bias=False, + ): + """Construct Lightweight 2-Dimensional Convolution layer.""" + super(LightweightConvolution2D, self).__init__() + + assert n_feat % wshare == 0 + self.wshare = wshare + self.use_kernel_mask = use_kernel_mask + self.dropout_rate = dropout_rate + self.kernel_size = kernel_size + self.padding_size = int(kernel_size / 2) + + # linear -> GLU -> lightconv -> linear + self.linear1 = nn.Linear(n_feat, n_feat * 2) + self.linear2 = nn.Linear(n_feat * 2, n_feat) + self.act = nn.GLU() + + # lightconv related + self.weight = nn.Parameter(torch.Tensor(self.wshare, 1, kernel_size).uniform_(0, 1)) + self.weight_f = nn.Parameter(torch.Tensor(1, 1, kernel_size).uniform_(0, 1)) + self.use_bias = use_bias + if self.use_bias: + self.bias = nn.Parameter(torch.Tensor(n_feat)) + + # mask of kernel + kernel_mask0 = torch.zeros(self.wshare, int(kernel_size / 2)) + kernel_mask1 = torch.ones(self.wshare, int(kernel_size / 2 + 1)) + self.kernel_mask = torch.cat((kernel_mask1, kernel_mask0), dim=-1).unsqueeze(1) + + def forward(self, query, key, value, mask): + """Forward of 'Lightweight 2-Dimensional Convolution'. + + This function takes query, key and value but uses only query. + This is just for compatibility with self-attention layer (attention.py) + + Args: + query (torch.Tensor): (batch, time1, d_model) input tensor + key (torch.Tensor): (batch, time2, d_model) NOT USED + value (torch.Tensor): (batch, time2, d_model) NOT USED + mask (torch.Tensor): (batch, time1, time2) mask + + Return: + x (torch.Tensor): (batch, time1, d_model) output + + """ + # linear -> GLU -> lightconv -> linear + x = query + B, T, C = x.size() + H = self.wshare + + # first liner layer + x = self.linear1(x) + + # GLU activation + x = self.act(x) + + # convolution along frequency axis + weight_f = F.softmax(self.weight_f, dim=-1) + weight_f = F.dropout(weight_f, self.dropout_rate, training=self.training) + weight_new = torch.zeros(B * T, 1, self.kernel_size, device=x.device, dtype=x.dtype).copy_( + weight_f + ) + xf = F.conv1d( + x.view(1, B * T, C), weight_new, padding=self.padding_size, groups=B * T + ).view(B, T, C) + + # lightconv + x = x.transpose(1, 2).contiguous().view(-1, H, T) # B x C x T + weight = F.dropout(self.weight, self.dropout_rate, training=self.training) + if self.use_kernel_mask: + self.kernel_mask = self.kernel_mask.to(x.device) + weight = weight.masked_fill(self.kernel_mask == 0.0, float("-inf")) + weight = F.softmax(weight, dim=-1) + x = F.conv1d(x, weight, padding=self.padding_size, groups=self.wshare).view(B, C, T) + if self.use_bias: + x = x + self.bias.view(1, -1, 1) + x = x.transpose(1, 2) # B x T x C + x = torch.cat((x, xf), -1) # B x T x Cx2 + + if mask is not None and not self.use_kernel_mask: + mask = mask.transpose(-1, -2) + x = x.masked_fill(mask == 0, 0.0) + + # second linear layer + x = self.linear2(x) + return x diff --git a/funasr/models/transformer/utils/mask.py b/funasr/models/transformer/utils/mask.py new file mode 100644 index 0000000000000000000000000000000000000000..e859e0e06c20fba953de8a6a0973ddc282b21b33 --- /dev/null +++ b/funasr/models/transformer/utils/mask.py @@ -0,0 +1,52 @@ +# Copyright 2019 Shigeki Karita +# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +"""Mask module.""" + +import torch + + +def subsequent_mask(size, device="cpu", dtype=torch.bool): + """Create mask for subsequent steps (size, size). + + :param int size: size of mask + :param str device: "cpu" or "cuda" or torch.Tensor.device + :param torch.dtype dtype: result dtype + :rtype: torch.Tensor + >>> subsequent_mask(3) + [[1, 0, 0], + [1, 1, 0], + [1, 1, 1]] + """ + ret = torch.ones(size, size, device=device, dtype=dtype) + return torch.tril(ret, out=ret) + + +def target_mask(ys_in_pad, ignore_id): + """Create mask for decoder self-attention. + + :param torch.Tensor ys_pad: batch of padded target sequences (B, Lmax) + :param int ignore_id: index of padding + :param torch.dtype dtype: result dtype + :rtype: torch.Tensor (B, Lmax, Lmax) + """ + ys_mask = ys_in_pad != ignore_id + m = subsequent_mask(ys_mask.size(-1), device=ys_mask.device).unsqueeze(0) + return ys_mask.unsqueeze(-2) & m + + +def vad_mask(size, vad_pos, device="cpu", dtype=torch.bool): + """Create mask for decoder self-attention. + + :param int size: size of mask + :param int vad_pos: index of vad index + :param str device: "cpu" or "cuda" or torch.Tensor.device + :param torch.dtype dtype: result dtype + :rtype: torch.Tensor (B, Lmax, Lmax) + """ + ret = torch.ones(size, size, device=device, dtype=dtype) + if vad_pos <= 0 or vad_pos >= size: + return ret + sub_corner = torch.zeros(vad_pos - 1, size - vad_pos, device=device, dtype=dtype) + ret[0 : vad_pos - 1, vad_pos:] = sub_corner + return ret diff --git a/funasr/models/transformer/utils/multi_layer_conv.py b/funasr/models/transformer/utils/multi_layer_conv.py new file mode 100644 index 0000000000000000000000000000000000000000..4893da167d406e399865531ee716bc199f3ab547 --- /dev/null +++ b/funasr/models/transformer/utils/multi_layer_conv.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +# Copyright 2019 Tomoki Hayashi +# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +"""Layer modules for FFT block in FastSpeech (Feed-forward Transformer).""" + +import torch + + +class MultiLayeredConv1d(torch.nn.Module): + """Multi-layered conv1d for Transformer block. + + This is a module of multi-leyered conv1d designed + to replace positionwise feed-forward network + in Transforner block, which is introduced in + `FastSpeech: Fast, Robust and Controllable Text to Speech`_. + + .. _`FastSpeech: Fast, Robust and Controllable Text to Speech`: + https://arxiv.org/pdf/1905.09263.pdf + + """ + + def __init__(self, in_chans, hidden_chans, kernel_size, dropout_rate): + """Initialize MultiLayeredConv1d module. + + Args: + in_chans (int): Number of input channels. + hidden_chans (int): Number of hidden channels. + kernel_size (int): Kernel size of conv1d. + dropout_rate (float): Dropout rate. + + """ + super(MultiLayeredConv1d, self).__init__() + self.w_1 = torch.nn.Conv1d( + in_chans, + hidden_chans, + kernel_size, + stride=1, + padding=(kernel_size - 1) // 2, + ) + self.w_2 = torch.nn.Conv1d( + hidden_chans, + in_chans, + kernel_size, + stride=1, + padding=(kernel_size - 1) // 2, + ) + self.dropout = torch.nn.Dropout(dropout_rate) + + def forward(self, x): + """Calculate forward propagation. + + Args: + x (torch.Tensor): Batch of input tensors (B, T, in_chans). + + Returns: + torch.Tensor: Batch of output tensors (B, T, hidden_chans). + + """ + x = torch.relu(self.w_1(x.transpose(-1, 1))).transpose(-1, 1) + return self.w_2(self.dropout(x).transpose(-1, 1)).transpose(-1, 1) + + +class FsmnFeedForward(torch.nn.Module): + """Position-wise feed forward for FSMN blocks. + + This is a module of multi-leyered conv1d designed + to replace position-wise feed-forward network + in FSMN block. + """ + + def __init__(self, in_chans, hidden_chans, out_chans, kernel_size, dropout_rate): + """Initialize FsmnFeedForward module. + + Args: + in_chans (int): Number of input channels. + hidden_chans (int): Number of hidden channels. + out_chans (int): Number of output channels. + kernel_size (int): Kernel size of conv1d. + dropout_rate (float): Dropout rate. + + """ + super(FsmnFeedForward, self).__init__() + self.w_1 = torch.nn.Conv1d( + in_chans, + hidden_chans, + kernel_size, + stride=1, + padding=(kernel_size - 1) // 2, + ) + self.w_2 = torch.nn.Conv1d( + hidden_chans, + out_chans, + kernel_size, + stride=1, + padding=(kernel_size - 1) // 2, + bias=False, + ) + self.norm = torch.nn.LayerNorm(hidden_chans) + self.dropout = torch.nn.Dropout(dropout_rate) + + def forward(self, x, ilens=None): + """Calculate forward propagation. + + Args: + x (torch.Tensor): Batch of input tensors (B, T, in_chans). + + Returns: + torch.Tensor: Batch of output tensors (B, T, out_chans). + + """ + x = torch.relu(self.w_1(x.transpose(-1, 1))).transpose(-1, 1) + return self.w_2(self.norm(self.dropout(x)).transpose(-1, 1)).transpose(-1, 1), ilens + + +class Conv1dLinear(torch.nn.Module): + """Conv1D + Linear for Transformer block. + + A variant of MultiLayeredConv1d, which replaces second conv-layer to linear. + + """ + + def __init__(self, in_chans, hidden_chans, kernel_size, dropout_rate): + """Initialize Conv1dLinear module. + + Args: + in_chans (int): Number of input channels. + hidden_chans (int): Number of hidden channels. + kernel_size (int): Kernel size of conv1d. + dropout_rate (float): Dropout rate. + + """ + super(Conv1dLinear, self).__init__() + self.w_1 = torch.nn.Conv1d( + in_chans, + hidden_chans, + kernel_size, + stride=1, + padding=(kernel_size - 1) // 2, + ) + self.w_2 = torch.nn.Linear(hidden_chans, in_chans) + self.dropout = torch.nn.Dropout(dropout_rate) + + def forward(self, x): + """Calculate forward propagation. + + Args: + x (torch.Tensor): Batch of input tensors (B, T, in_chans). + + Returns: + torch.Tensor: Batch of output tensors (B, T, hidden_chans). + + """ + x = torch.relu(self.w_1(x.transpose(-1, 1))).transpose(-1, 1) + return self.w_2(self.dropout(x)) diff --git a/funasr/models/transformer/utils/nets_utils.py b/funasr/models/transformer/utils/nets_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..29d23ee59c75068d2bafc1f09b76753636e0be7f --- /dev/null +++ b/funasr/models/transformer/utils/nets_utils.py @@ -0,0 +1,734 @@ +# -*- coding: utf-8 -*- + +"""Network related utility tools.""" + +import logging +from typing import Dict, List, Tuple + +import numpy as np +import torch + + +def to_device(m, x): + """Send tensor into the device of the module. + + Args: + m (torch.nn.Module): Torch module. + x (Tensor): Torch tensor. + + Returns: + Tensor: Torch tensor located in the same place as torch module. + + """ + if isinstance(m, torch.nn.Module): + device = next(m.parameters()).device + elif isinstance(m, torch.Tensor): + device = m.device + else: + raise TypeError("Expected torch.nn.Module or torch.tensor, " f"bot got: {type(m)}") + return x.to(device) + + +def pad_list(xs, pad_value): + """Perform padding for the list of tensors. + + Args: + xs (List): List of Tensors [(T_1, `*`), (T_2, `*`), ..., (T_B, `*`)]. + pad_value (float): Value for padding. + + Returns: + Tensor: Padded tensor (B, Tmax, `*`). + + Examples: + >>> x = [torch.ones(4), torch.ones(2), torch.ones(1)] + >>> x + [tensor([1., 1., 1., 1.]), tensor([1., 1.]), tensor([1.])] + >>> pad_list(x, 0) + tensor([[1., 1., 1., 1.], + [1., 1., 0., 0.], + [1., 0., 0., 0.]]) + + """ + n_batch = len(xs) + max_len = max(x.size(0) for x in xs) + pad = xs[0].new(n_batch, max_len, *xs[0].size()[1:]).fill_(pad_value) + + for i in range(n_batch): + pad[i, : xs[i].size(0)] = xs[i] + + return pad + + +def pad_list_all_dim(xs, pad_value): + """Perform padding for the list of tensors. + + Args: + xs (List): List of Tensors [(T_1, `*`), (T_2, `*`), ..., (T_B, `*`)]. + pad_value (float): Value for padding. + + Returns: + Tensor: Padded tensor (B, Tmax, `*`). + + Examples: + >>> x = [torch.ones(4), torch.ones(2), torch.ones(1)] + >>> x + [tensor([1., 1., 1., 1.]), tensor([1., 1.]), tensor([1.])] + >>> pad_list(x, 0) + tensor([[1., 1., 1., 1.], + [1., 1., 0., 0.], + [1., 0., 0., 0.]]) + + """ + n_batch = len(xs) + num_dim = len(xs[0].shape) + max_len_all_dim = [] + for i in range(num_dim): + max_len_all_dim.append(max(x.size(i) for x in xs)) + pad = xs[0].new(n_batch, *max_len_all_dim).fill_(pad_value) + + for i in range(n_batch): + if num_dim == 1: + pad[i, : xs[i].size(0)] = xs[i] + elif num_dim == 2: + pad[i, : xs[i].size(0), : xs[i].size(1)] = xs[i] + elif num_dim == 3: + pad[i, : xs[i].size(0), : xs[i].size(1), : xs[i].size(2)] = xs[i] + else: + raise ValueError( + "pad_list_all_dim only support 1-D, 2-D and 3-D tensors, not {}-D.".format(num_dim) + ) + + return pad + + +def make_pad_mask(lengths, xs=None, length_dim=-1, maxlen=None): + """Make mask tensor containing indices of padded part. + + Args: + lengths (LongTensor or List): Batch of lengths (B,). + xs (Tensor, optional): The reference tensor. + If set, masks will be the same shape as this tensor. + length_dim (int, optional): Dimension indicator of the above tensor. + See the example. + + Returns: + Tensor: Mask tensor containing indices of padded part. + dtype=torch.uint8 in PyTorch 1.2- + dtype=torch.bool in PyTorch 1.2+ (including 1.2) + + Examples: + With only lengths. + + >>> lengths = [5, 3, 2] + >>> make_pad_mask(lengths) + masks = [[0, 0, 0, 0 ,0], + [0, 0, 0, 1, 1], + [0, 0, 1, 1, 1]] + + With the reference tensor. + + >>> xs = torch.zeros((3, 2, 4)) + >>> make_pad_mask(lengths, xs) + tensor([[[0, 0, 0, 0], + [0, 0, 0, 0]], + [[0, 0, 0, 1], + [0, 0, 0, 1]], + [[0, 0, 1, 1], + [0, 0, 1, 1]]], dtype=torch.uint8) + >>> xs = torch.zeros((3, 2, 6)) + >>> make_pad_mask(lengths, xs) + tensor([[[0, 0, 0, 0, 0, 1], + [0, 0, 0, 0, 0, 1]], + [[0, 0, 0, 1, 1, 1], + [0, 0, 0, 1, 1, 1]], + [[0, 0, 1, 1, 1, 1], + [0, 0, 1, 1, 1, 1]]], dtype=torch.uint8) + + With the reference tensor and dimension indicator. + + >>> xs = torch.zeros((3, 6, 6)) + >>> make_pad_mask(lengths, xs, 1) + tensor([[[0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [1, 1, 1, 1, 1, 1]], + [[0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1]], + [[0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1]]], dtype=torch.uint8) + >>> make_pad_mask(lengths, xs, 2) + tensor([[[0, 0, 0, 0, 0, 1], + [0, 0, 0, 0, 0, 1], + [0, 0, 0, 0, 0, 1], + [0, 0, 0, 0, 0, 1], + [0, 0, 0, 0, 0, 1], + [0, 0, 0, 0, 0, 1]], + [[0, 0, 0, 1, 1, 1], + [0, 0, 0, 1, 1, 1], + [0, 0, 0, 1, 1, 1], + [0, 0, 0, 1, 1, 1], + [0, 0, 0, 1, 1, 1], + [0, 0, 0, 1, 1, 1]], + [[0, 0, 1, 1, 1, 1], + [0, 0, 1, 1, 1, 1], + [0, 0, 1, 1, 1, 1], + [0, 0, 1, 1, 1, 1], + [0, 0, 1, 1, 1, 1], + [0, 0, 1, 1, 1, 1]]], dtype=torch.uint8) + + """ + if length_dim == 0: + raise ValueError("length_dim cannot be 0: {}".format(length_dim)) + + if not isinstance(lengths, list): + lengths = lengths.tolist() + bs = int(len(lengths)) + if maxlen is None: + if xs is None: + maxlen = int(max(lengths)) + else: + maxlen = xs.size(length_dim) + else: + assert xs is None + assert maxlen >= int(max(lengths)) + + seq_range = torch.arange(0, maxlen, dtype=torch.int64) + seq_range_expand = seq_range.unsqueeze(0).expand(bs, maxlen) + seq_length_expand = seq_range_expand.new(lengths).unsqueeze(-1) + mask = seq_range_expand >= seq_length_expand + + if xs is not None: + assert xs.size(0) == bs, (xs.size(0), bs) + + if length_dim < 0: + length_dim = xs.dim() + length_dim + # ind = (:, None, ..., None, :, , None, ..., None) + ind = tuple(slice(None) if i in (0, length_dim) else None for i in range(xs.dim())) + mask = mask[ind].expand_as(xs).to(xs.device) + return mask + + +def make_non_pad_mask(lengths, xs=None, length_dim=-1): + """Make mask tensor containing indices of non-padded part. + + Args: + lengths (LongTensor or List): Batch of lengths (B,). + xs (Tensor, optional): The reference tensor. + If set, masks will be the same shape as this tensor. + length_dim (int, optional): Dimension indicator of the above tensor. + See the example. + + Returns: + ByteTensor: mask tensor containing indices of padded part. + dtype=torch.uint8 in PyTorch 1.2- + dtype=torch.bool in PyTorch 1.2+ (including 1.2) + + Examples: + With only lengths. + + >>> lengths = [5, 3, 2] + >>> make_non_pad_mask(lengths) + masks = [[1, 1, 1, 1 ,1], + [1, 1, 1, 0, 0], + [1, 1, 0, 0, 0]] + + With the reference tensor. + + >>> xs = torch.zeros((3, 2, 4)) + >>> make_non_pad_mask(lengths, xs) + tensor([[[1, 1, 1, 1], + [1, 1, 1, 1]], + [[1, 1, 1, 0], + [1, 1, 1, 0]], + [[1, 1, 0, 0], + [1, 1, 0, 0]]], dtype=torch.uint8) + >>> xs = torch.zeros((3, 2, 6)) + >>> make_non_pad_mask(lengths, xs) + tensor([[[1, 1, 1, 1, 1, 0], + [1, 1, 1, 1, 1, 0]], + [[1, 1, 1, 0, 0, 0], + [1, 1, 1, 0, 0, 0]], + [[1, 1, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0]]], dtype=torch.uint8) + + With the reference tensor and dimension indicator. + + >>> xs = torch.zeros((3, 6, 6)) + >>> make_non_pad_mask(lengths, xs, 1) + tensor([[[1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1], + [0, 0, 0, 0, 0, 0]], + [[1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0]], + [[1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0]]], dtype=torch.uint8) + >>> make_non_pad_mask(lengths, xs, 2) + tensor([[[1, 1, 1, 1, 1, 0], + [1, 1, 1, 1, 1, 0], + [1, 1, 1, 1, 1, 0], + [1, 1, 1, 1, 1, 0], + [1, 1, 1, 1, 1, 0], + [1, 1, 1, 1, 1, 0]], + [[1, 1, 1, 0, 0, 0], + [1, 1, 1, 0, 0, 0], + [1, 1, 1, 0, 0, 0], + [1, 1, 1, 0, 0, 0], + [1, 1, 1, 0, 0, 0], + [1, 1, 1, 0, 0, 0]], + [[1, 1, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0]]], dtype=torch.uint8) + + """ + return ~make_pad_mask(lengths, xs, length_dim) + + +def mask_by_length(xs, lengths, fill=0): + """Mask tensor according to length. + + Args: + xs (Tensor): Batch of input tensor (B, `*`). + lengths (LongTensor or List): Batch of lengths (B,). + fill (int or float): Value to fill masked part. + + Returns: + Tensor: Batch of masked input tensor (B, `*`). + + Examples: + >>> x = torch.arange(5).repeat(3, 1) + 1 + >>> x + tensor([[1, 2, 3, 4, 5], + [1, 2, 3, 4, 5], + [1, 2, 3, 4, 5]]) + >>> lengths = [5, 3, 2] + >>> mask_by_length(x, lengths) + tensor([[1, 2, 3, 4, 5], + [1, 2, 3, 0, 0], + [1, 2, 0, 0, 0]]) + + """ + assert xs.size(0) == len(lengths) + ret = xs.data.new(*xs.size()).fill_(fill) + for i, l in enumerate(lengths): + ret[i, :l] = xs[i, :l] + return ret + + +def to_torch_tensor(x): + """Change to torch.Tensor or ComplexTensor from numpy.ndarray. + + Args: + x: Inputs. It should be one of numpy.ndarray, Tensor, ComplexTensor, and dict. + + Returns: + Tensor or ComplexTensor: Type converted inputs. + + Examples: + >>> xs = np.ones(3, dtype=np.float32) + >>> xs = to_torch_tensor(xs) + tensor([1., 1., 1.]) + >>> xs = torch.ones(3, 4, 5) + >>> assert to_torch_tensor(xs) is xs + >>> xs = {'real': xs, 'imag': xs} + >>> to_torch_tensor(xs) + ComplexTensor( + Real: + tensor([1., 1., 1.]) + Imag; + tensor([1., 1., 1.]) + ) + + """ + # If numpy, change to torch tensor + if isinstance(x, np.ndarray): + if x.dtype.kind == "c": + # Dynamically importing because torch_complex requires python3 + from torch_complex.tensor import ComplexTensor + + return ComplexTensor(x) + else: + return torch.from_numpy(x) + + # If {'real': ..., 'imag': ...}, convert to ComplexTensor + elif isinstance(x, dict): + # Dynamically importing because torch_complex requires python3 + from torch_complex.tensor import ComplexTensor + + if "real" not in x or "imag" not in x: + raise ValueError("has 'real' and 'imag' keys: {}".format(list(x))) + # Relative importing because of using python3 syntax + return ComplexTensor(x["real"], x["imag"]) + + # If torch.Tensor, as it is + elif isinstance(x, torch.Tensor): + return x + + else: + error = ( + "x must be numpy.ndarray, torch.Tensor or a dict like " + "{{'real': torch.Tensor, 'imag': torch.Tensor}}, " + "but got {}".format(type(x)) + ) + try: + from torch_complex.tensor import ComplexTensor + except Exception: + # If PY2 + raise ValueError(error) + else: + # If PY3 + if isinstance(x, ComplexTensor): + return x + else: + raise ValueError(error) + + +def get_subsample(train_args, mode, arch): + """Parse the subsampling factors from the args for the specified `mode` and `arch`. + + Args: + train_args: argument Namespace containing options. + mode: one of ('asr', 'mt', 'st') + arch: one of ('rnn', 'rnn-t', 'rnn_mix', 'rnn_mulenc', 'transformer') + + Returns: + np.ndarray / List[np.ndarray]: subsampling factors. + """ + if arch == "transformer": + return np.array([1]) + + elif mode == "mt" and arch == "rnn": + # +1 means input (+1) and layers outputs (train_args.elayer) + subsample = np.ones(train_args.elayers + 1, dtype=np.int32) + logging.warning("Subsampling is not performed for machine translation.") + logging.info("subsample: " + " ".join([str(x) for x in subsample])) + return subsample + + elif ( + (mode == "asr" and arch in ("rnn", "rnn-t")) + or (mode == "mt" and arch == "rnn") + or (mode == "st" and arch == "rnn") + ): + subsample = np.ones(train_args.elayers + 1, dtype=np.int32) + if train_args.etype.endswith("p") and not train_args.etype.startswith("vgg"): + ss = train_args.subsample.split("_") + for j in range(min(train_args.elayers + 1, len(ss))): + subsample[j] = int(ss[j]) + else: + logging.warning( + "Subsampling is not performed for vgg*. " + "It is performed in max pooling layers at CNN." + ) + logging.info("subsample: " + " ".join([str(x) for x in subsample])) + return subsample + + elif mode == "asr" and arch == "rnn_mix": + subsample = np.ones(train_args.elayers_sd + train_args.elayers + 1, dtype=np.int32) + if train_args.etype.endswith("p") and not train_args.etype.startswith("vgg"): + ss = train_args.subsample.split("_") + for j in range(min(train_args.elayers_sd + train_args.elayers + 1, len(ss))): + subsample[j] = int(ss[j]) + else: + logging.warning( + "Subsampling is not performed for vgg*. " + "It is performed in max pooling layers at CNN." + ) + logging.info("subsample: " + " ".join([str(x) for x in subsample])) + return subsample + + elif mode == "asr" and arch == "rnn_mulenc": + subsample_list = [] + for idx in range(train_args.num_encs): + subsample = np.ones(train_args.elayers[idx] + 1, dtype=np.int32) + if train_args.etype[idx].endswith("p") and not train_args.etype[idx].startswith("vgg"): + ss = train_args.subsample[idx].split("_") + for j in range(min(train_args.elayers[idx] + 1, len(ss))): + subsample[j] = int(ss[j]) + else: + logging.warning( + "Encoder %d: Subsampling is not performed for vgg*. " + "It is performed in max pooling layers at CNN.", + idx + 1, + ) + logging.info("subsample: " + " ".join([str(x) for x in subsample])) + subsample_list.append(subsample) + return subsample_list + + else: + raise ValueError("Invalid options: mode={}, arch={}".format(mode, arch)) + + +def rename_state_dict(old_prefix: str, new_prefix: str, state_dict: Dict[str, torch.Tensor]): + """Replace keys of old prefix with new prefix in state dict.""" + # need this list not to break the dict iterator + old_keys = [k for k in state_dict if k.startswith(old_prefix)] + if len(old_keys) > 0: + logging.warning(f"Rename: {old_prefix} -> {new_prefix}") + for k in old_keys: + v = state_dict.pop(k) + new_k = k.replace(old_prefix, new_prefix) + state_dict[new_k] = v + + +class Swish(torch.nn.Module): + """Swish activation definition. + + Swish(x) = (beta * x) * sigmoid(x) + where beta = 1 defines standard Swish activation. + + References: + https://arxiv.org/abs/2108.12943 / https://arxiv.org/abs/1710.05941v1. + E-swish variant: https://arxiv.org/abs/1801.07145. + + Args: + beta: Beta parameter for E-Swish. + (beta >= 1. If beta < 1, use standard Swish). + use_builtin: Whether to use PyTorch function if available. + + """ + + def __init__(self, beta: float = 1.0, use_builtin: bool = False) -> None: + super().__init__() + + self.beta = beta + + if beta > 1: + self.swish = lambda x: (self.beta * x) * torch.sigmoid(x) + else: + if use_builtin: + self.swish = torch.nn.SiLU() + else: + self.swish = lambda x: x * torch.sigmoid(x) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward computation.""" + return self.swish(x) + + +def get_activation(act): + """Return activation function.""" + + activation_funcs = { + "hardtanh": torch.nn.Hardtanh, + "tanh": torch.nn.Tanh, + "relu": torch.nn.ReLU, + "selu": torch.nn.SELU, + "swish": Swish, + } + + return activation_funcs[act]() + + +class TooShortUttError(Exception): + """Raised when the utt is too short for subsampling. + + Args: + message: Error message to display. + actual_size: The size that cannot pass the subsampling. + limit: The size limit for subsampling. + + """ + + def __init__(self, message: str, actual_size: int, limit: int) -> None: + """Construct a TooShortUttError module.""" + super().__init__(message) + + self.actual_size = actual_size + self.limit = limit + + +def check_short_utt(sub_factor: int, size: int) -> Tuple[bool, int]: + """Check if the input is too short for subsampling. + + Args: + sub_factor: Subsampling factor for Conv2DSubsampling. + size: Input size. + + Returns: + : Whether an error should be sent. + : Size limit for specified subsampling factor. + + """ + if sub_factor == 2 and size < 3: + return True, 7 + elif sub_factor == 4 and size < 7: + return True, 7 + elif sub_factor == 6 and size < 11: + return True, 11 + + return False, -1 + + +def sub_factor_to_params(sub_factor: int, input_size: int) -> Tuple[int, int, int]: + """Get conv2D second layer parameters for given subsampling factor. + + Args: + sub_factor: Subsampling factor (1/X). + input_size: Input size. + + Returns: + : Kernel size for second convolution. + : Stride for second convolution. + : Conv2DSubsampling output size. + + """ + if sub_factor == 2: + return 3, 1, (((input_size - 1) // 2 - 2)) + elif sub_factor == 4: + return 3, 2, (((input_size - 1) // 2 - 1) // 2) + elif sub_factor == 6: + return 5, 3, (((input_size - 1) // 2 - 2) // 3) + else: + raise ValueError("subsampling_factor parameter should be set to either 2, 4 or 6.") + + +def make_chunk_mask( + size: int, + chunk_size: int, + left_chunk_size: int = 0, + device: torch.device = None, +) -> torch.Tensor: + """Create chunk mask for the subsequent steps (size, size). + + Reference: https://github.com/k2-fsa/icefall/blob/master/icefall/utils.py + + Args: + size: Size of the source mask. + chunk_size: Number of frames in chunk. + left_chunk_size: Size of the left context in chunks (0 means full context). + device: Device for the mask tensor. + + Returns: + mask: Chunk mask. (size, size) + + """ + mask = torch.zeros(size, size, device=device, dtype=torch.bool) + + for i in range(size): + if left_chunk_size < 0: + start = 0 + else: + start = max((i // chunk_size - left_chunk_size) * chunk_size, 0) + + end = min((i // chunk_size + 1) * chunk_size, size) + mask[i, start:end] = True + + return ~mask + + +def make_source_mask(lengths: torch.Tensor) -> torch.Tensor: + """Create source mask for given lengths. + + Reference: https://github.com/k2-fsa/icefall/blob/master/icefall/utils.py + + Args: + lengths: Sequence lengths. (B,) + + Returns: + : Mask for the sequence lengths. (B, max_len) + + """ + max_len = lengths.max() + batch_size = lengths.size(0) + + expanded_lengths = torch.arange(max_len).expand(batch_size, max_len).to(lengths) + + return expanded_lengths >= lengths.unsqueeze(1) + + +def get_transducer_task_io( + labels: torch.Tensor, + encoder_out_lens: torch.Tensor, + ignore_id: int = -1, + blank_id: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Get Transducer loss I/O. + + Args: + labels: Label ID sequences. (B, L) + encoder_out_lens: Encoder output lengths. (B,) + ignore_id: Padding symbol ID. + blank_id: Blank symbol ID. + + Returns: + decoder_in: Decoder inputs. (B, U) + target: Target label ID sequences. (B, U) + t_len: Time lengths. (B,) + u_len: Label lengths. (B,) + + """ + + def pad_list(labels: List[torch.Tensor], padding_value: int = 0): + """Create padded batch of labels from a list of labels sequences. + + Args: + labels: Labels sequences. [B x (?)] + padding_value: Padding value. + + Returns: + labels: Batch of padded labels sequences. (B,) + + """ + batch_size = len(labels) + + padded = ( + labels[0] + .new(batch_size, max(x.size(0) for x in labels), *labels[0].size()[1:]) + .fill_(padding_value) + ) + + for i in range(batch_size): + padded[i, : labels[i].size(0)] = labels[i] + + return padded + + device = labels.device + + labels_unpad = [y[y != ignore_id] for y in labels] + blank = labels[0].new([blank_id]) + + decoder_in = pad_list( + [torch.cat([blank, label], dim=0) for label in labels_unpad], blank_id + ).to(device) + + target = pad_list(labels_unpad, blank_id).type(torch.int32).to(device) + + encoder_out_lens = list(map(int, encoder_out_lens)) + t_len = torch.IntTensor(encoder_out_lens).to(device) + + u_len = torch.IntTensor([y.size(0) for y in labels_unpad]).to(device) + + return decoder_in, target, t_len, u_len + + +def pad_to_len(t: torch.Tensor, pad_len: int, dim: int): + """Pad the tensor `t` at `dim` to the length `pad_len` with right padding zeros.""" + if t.size(dim) == pad_len: + return t + else: + pad_size = list(t.shape) + pad_size[dim] = pad_len - t.size(dim) + return torch.cat([t, torch.zeros(*pad_size, dtype=t.dtype, device=t.device)], dim=dim) diff --git a/funasr/models/transformer/utils/repeat.py b/funasr/models/transformer/utils/repeat.py new file mode 100644 index 0000000000000000000000000000000000000000..a44c1a0103a1877678a7df462af742bb1cd32de1 --- /dev/null +++ b/funasr/models/transformer/utils/repeat.py @@ -0,0 +1,137 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +# Copyright 2019 Shigeki Karita +# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +"""Repeat the same layer definition.""" + +from typing import Dict, List, Optional +from funasr.models.transformer.layer_norm import LayerNorm +import torch + + +class MultiSequential(torch.nn.Sequential): + """Multi-input multi-output torch.nn.Sequential.""" + + def __init__(self, *args, layer_drop_rate=0.0): + """Initialize MultiSequential with layer_drop. + + Args: + layer_drop_rate (float): Probability of dropping out each fn (layer). + + """ + super(MultiSequential, self).__init__(*args) + self.layer_drop_rate = layer_drop_rate + + def forward(self, *args): + """Repeat.""" + _probs = torch.empty(len(self)).uniform_() + for idx, m in enumerate(self): + if not self.training or (_probs[idx] >= self.layer_drop_rate): + args = m(*args) + return args + + +def repeat(N, fn, layer_drop_rate=0.0): + """Repeat module N times. + + Args: + N (int): Number of repeat time. + fn (Callable): Function to generate module. + layer_drop_rate (float): Probability of dropping out each fn (layer). + + Returns: + MultiSequential: Repeated model instance. + + """ + return MultiSequential(*[fn(n) for n in range(N)], layer_drop_rate=layer_drop_rate) + + +class MultiBlocks(torch.nn.Module): + """MultiBlocks definition. + Args: + block_list: Individual blocks of the encoder architecture. + output_size: Architecture output size. + norm_class: Normalization module class. + norm_args: Normalization module arguments. + """ + + def __init__( + self, + block_list: List[torch.nn.Module], + output_size: int, + norm_class: torch.nn.Module = LayerNorm, + ) -> None: + """Construct a MultiBlocks object.""" + super().__init__() + + self.blocks = torch.nn.ModuleList(block_list) + self.norm_blocks = norm_class(output_size) + + self.num_blocks = len(block_list) + + def reset_streaming_cache(self, left_context: int, device: torch.device) -> None: + """Initialize/Reset encoder streaming cache. + Args: + left_context: Number of left frames during chunk-by-chunk inference. + device: Device to use for cache tensor. + """ + for idx in range(self.num_blocks): + self.blocks[idx].reset_streaming_cache(left_context, device) + + def forward( + self, + x: torch.Tensor, + pos_enc: torch.Tensor, + mask: torch.Tensor, + chunk_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Forward each block of the encoder architecture. + Args: + x: MultiBlocks input sequences. (B, T, D_block_1) + pos_enc: Positional embedding sequences. + mask: Source mask. (B, T) + chunk_mask: Chunk mask. (T_2, T_2) + Returns: + x: Output sequences. (B, T, D_block_N) + """ + for block_index, block in enumerate(self.blocks): + x, mask, pos_enc = block(x, pos_enc, mask, chunk_mask=chunk_mask) + + x = self.norm_blocks(x) + + return x + + def chunk_forward( + self, + x: torch.Tensor, + pos_enc: torch.Tensor, + mask: torch.Tensor, + chunk_size: int = 0, + left_context: int = 0, + right_context: int = 0, + ) -> torch.Tensor: + """Forward each block of the encoder architecture. + Args: + x: MultiBlocks input sequences. (B, T, D_block_1) + pos_enc: Positional embedding sequences. (B, 2 * (T - 1), D_att) + mask: Source mask. (B, T_2) + left_context: Number of frames in left context. + right_context: Number of frames in right context. + Returns: + x: MultiBlocks output sequences. (B, T, D_block_N) + """ + for block_idx, block in enumerate(self.blocks): + x, pos_enc = block.chunk_forward( + x, + pos_enc, + mask, + chunk_size=chunk_size, + left_context=left_context, + right_context=right_context, + ) + + x = self.norm_blocks(x) + + return x diff --git a/funasr/models/transformer/utils/subsampling.py b/funasr/models/transformer/utils/subsampling.py new file mode 100644 index 0000000000000000000000000000000000000000..b18bd3132ed2cd45db7118217c4968c454a19024 --- /dev/null +++ b/funasr/models/transformer/utils/subsampling.py @@ -0,0 +1,629 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +# Copyright 2019 Shigeki Karita +# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +"""Subsampling layer definition.""" +import numpy as np +import torch +import torch.nn.functional as F +from funasr.models.transformer.embedding import PositionalEncoding +import logging +from funasr.models.scama.utils import sequence_mask +from funasr.models.transformer.utils.nets_utils import sub_factor_to_params, pad_to_len +from typing import Optional, Tuple, Union +import math + + +class TooShortUttError(Exception): + """Raised when the utt is too short for subsampling. + + Args: + message (str): Message for error catch + actual_size (int): the short size that cannot pass the subsampling + limit (int): the limit size for subsampling + + """ + + def __init__(self, message, actual_size, limit): + """Construct a TooShortUttError for error handler.""" + super().__init__(message) + self.actual_size = actual_size + self.limit = limit + + +def check_short_utt(ins, size): + """Check if the utterance is too short for subsampling.""" + if isinstance(ins, Conv2dSubsampling2) and size < 3: + return True, 3 + if isinstance(ins, Conv2dSubsampling) and size < 7: + return True, 7 + if isinstance(ins, Conv2dSubsampling6) and size < 11: + return True, 11 + if isinstance(ins, Conv2dSubsampling8) and size < 15: + return True, 15 + return False, -1 + + +class Conv2dSubsampling(torch.nn.Module): + """Convolutional 2D subsampling (to 1/4 length). + + Args: + idim (int): Input dimension. + odim (int): Output dimension. + dropout_rate (float): Dropout rate. + pos_enc (torch.nn.Module): Custom position encoding layer. + + """ + + def __init__(self, idim, odim, dropout_rate, pos_enc=None): + """Construct an Conv2dSubsampling object.""" + super(Conv2dSubsampling, self).__init__() + self.conv = torch.nn.Sequential( + torch.nn.Conv2d(1, odim, 3, 2), + torch.nn.ReLU(), + torch.nn.Conv2d(odim, odim, 3, 2), + torch.nn.ReLU(), + ) + self.out = torch.nn.Sequential( + torch.nn.Linear(odim * (((idim - 1) // 2 - 1) // 2), odim), + pos_enc if pos_enc is not None else PositionalEncoding(odim, dropout_rate), + ) + + def forward(self, x, x_mask): + """Subsample x. + + Args: + x (torch.Tensor): Input tensor (#batch, time, idim). + x_mask (torch.Tensor): Input mask (#batch, 1, time). + + Returns: + torch.Tensor: Subsampled tensor (#batch, time', odim), + where time' = time // 4. + torch.Tensor: Subsampled mask (#batch, 1, time'), + where time' = time // 4. + + """ + x = x.unsqueeze(1) # (b, c, t, f) + x = self.conv(x) + b, c, t, f = x.size() + x = self.out(x.transpose(1, 2).contiguous().view(b, t, c * f)) + if x_mask is None: + return x, None + return x, x_mask[:, :, :-2:2][:, :, :-2:2] + + def __getitem__(self, key): + """Get item. + + When reset_parameters() is called, if use_scaled_pos_enc is used, + return the positioning encoding. + + """ + if key != -1: + raise NotImplementedError("Support only `-1` (for `reset_parameters`).") + return self.out[key] + + +class Conv2dSubsamplingPad(torch.nn.Module): + """Convolutional 2D subsampling (to 1/4 length). + + Args: + idim (int): Input dimension. + odim (int): Output dimension. + dropout_rate (float): Dropout rate. + pos_enc (torch.nn.Module): Custom position encoding layer. + + """ + + def __init__(self, idim, odim, dropout_rate, pos_enc=None): + """Construct an Conv2dSubsampling object.""" + super(Conv2dSubsamplingPad, self).__init__() + self.conv = torch.nn.Sequential( + torch.nn.Conv2d(1, odim, 3, 2, padding=(0, 0)), + torch.nn.ReLU(), + torch.nn.Conv2d(odim, odim, 3, 2, padding=(0, 0)), + torch.nn.ReLU(), + ) + self.out = torch.nn.Sequential( + torch.nn.Linear(odim * (((idim - 1) // 2 - 1) // 2), odim), + pos_enc if pos_enc is not None else PositionalEncoding(odim, dropout_rate), + ) + self.pad_fn = torch.nn.ConstantPad1d((0, 4), 0.0) + + def forward(self, x, x_mask): + """Subsample x. + + Args: + x (torch.Tensor): Input tensor (#batch, time, idim). + x_mask (torch.Tensor): Input mask (#batch, 1, time). + + Returns: + torch.Tensor: Subsampled tensor (#batch, time', odim), + where time' = time // 4. + torch.Tensor: Subsampled mask (#batch, 1, time'), + where time' = time // 4. + + """ + x = x.transpose(1, 2) + x = self.pad_fn(x) + x = x.transpose(1, 2) + x = x.unsqueeze(1) # (b, c, t, f) + x = self.conv(x) + b, c, t, f = x.size() + x = self.out(x.transpose(1, 2).contiguous().view(b, t, c * f)) + if x_mask is None: + return x, None + x_len = torch.sum(x_mask[:, 0, :], dim=-1) + x_len = (x_len - 1) // 2 + 1 + x_len = (x_len - 1) // 2 + 1 + mask = sequence_mask(x_len, None, x_len.dtype, x[0].device) + return x, mask[:, None, :] + + def __getitem__(self, key): + """Get item. + + When reset_parameters() is called, if use_scaled_pos_enc is used, + return the positioning encoding. + + """ + if key != -1: + raise NotImplementedError("Support only `-1` (for `reset_parameters`).") + return self.out[key] + + +class Conv2dSubsampling2(torch.nn.Module): + """Convolutional 2D subsampling (to 1/2 length). + + Args: + idim (int): Input dimension. + odim (int): Output dimension. + dropout_rate (float): Dropout rate. + pos_enc (torch.nn.Module): Custom position encoding layer. + + """ + + def __init__(self, idim, odim, dropout_rate, pos_enc=None): + """Construct an Conv2dSubsampling2 object.""" + super(Conv2dSubsampling2, self).__init__() + self.conv = torch.nn.Sequential( + torch.nn.Conv2d(1, odim, 3, 2), + torch.nn.ReLU(), + torch.nn.Conv2d(odim, odim, 3, 1), + torch.nn.ReLU(), + ) + self.out = torch.nn.Sequential( + torch.nn.Linear(odim * (((idim - 1) // 2 - 2)), odim), + pos_enc if pos_enc is not None else PositionalEncoding(odim, dropout_rate), + ) + + def forward(self, x, x_mask): + """Subsample x. + + Args: + x (torch.Tensor): Input tensor (#batch, time, idim). + x_mask (torch.Tensor): Input mask (#batch, 1, time). + + Returns: + torch.Tensor: Subsampled tensor (#batch, time', odim), + where time' = time // 2. + torch.Tensor: Subsampled mask (#batch, 1, time'), + where time' = time // 2. + + """ + x = x.unsqueeze(1) # (b, c, t, f) + x = self.conv(x) + b, c, t, f = x.size() + x = self.out(x.transpose(1, 2).contiguous().view(b, t, c * f)) + if x_mask is None: + return x, None + return x, x_mask[:, :, :-2:2][:, :, :-2:1] + + def __getitem__(self, key): + """Get item. + + When reset_parameters() is called, if use_scaled_pos_enc is used, + return the positioning encoding. + + """ + if key != -1: + raise NotImplementedError("Support only `-1` (for `reset_parameters`).") + return self.out[key] + + +class Conv2dSubsampling6(torch.nn.Module): + """Convolutional 2D subsampling (to 1/6 length). + + Args: + idim (int): Input dimension. + odim (int): Output dimension. + dropout_rate (float): Dropout rate. + pos_enc (torch.nn.Module): Custom position encoding layer. + + """ + + def __init__(self, idim, odim, dropout_rate, pos_enc=None): + """Construct an Conv2dSubsampling6 object.""" + super(Conv2dSubsampling6, self).__init__() + self.conv = torch.nn.Sequential( + torch.nn.Conv2d(1, odim, 3, 2), + torch.nn.ReLU(), + torch.nn.Conv2d(odim, odim, 5, 3), + torch.nn.ReLU(), + ) + self.out = torch.nn.Sequential( + torch.nn.Linear(odim * (((idim - 1) // 2 - 2) // 3), odim), + pos_enc if pos_enc is not None else PositionalEncoding(odim, dropout_rate), + ) + + def forward(self, x, x_mask): + """Subsample x. + + Args: + x (torch.Tensor): Input tensor (#batch, time, idim). + x_mask (torch.Tensor): Input mask (#batch, 1, time). + + Returns: + torch.Tensor: Subsampled tensor (#batch, time', odim), + where time' = time // 6. + torch.Tensor: Subsampled mask (#batch, 1, time'), + where time' = time // 6. + + """ + x = x.unsqueeze(1) # (b, c, t, f) + x = self.conv(x) + b, c, t, f = x.size() + x = self.out(x.transpose(1, 2).contiguous().view(b, t, c * f)) + if x_mask is None: + return x, None + return x, x_mask[:, :, :-2:2][:, :, :-4:3] + + +class Conv2dSubsampling8(torch.nn.Module): + """Convolutional 2D subsampling (to 1/8 length). + + Args: + idim (int): Input dimension. + odim (int): Output dimension. + dropout_rate (float): Dropout rate. + pos_enc (torch.nn.Module): Custom position encoding layer. + + """ + + def __init__(self, idim, odim, dropout_rate, pos_enc=None): + """Construct an Conv2dSubsampling8 object.""" + super(Conv2dSubsampling8, self).__init__() + self.conv = torch.nn.Sequential( + torch.nn.Conv2d(1, odim, 3, 2), + torch.nn.ReLU(), + torch.nn.Conv2d(odim, odim, 3, 2), + torch.nn.ReLU(), + torch.nn.Conv2d(odim, odim, 3, 2), + torch.nn.ReLU(), + ) + self.out = torch.nn.Sequential( + torch.nn.Linear(odim * ((((idim - 1) // 2 - 1) // 2 - 1) // 2), odim), + pos_enc if pos_enc is not None else PositionalEncoding(odim, dropout_rate), + ) + + def forward(self, x, x_mask): + """Subsample x. + + Args: + x (torch.Tensor): Input tensor (#batch, time, idim). + x_mask (torch.Tensor): Input mask (#batch, 1, time). + + Returns: + torch.Tensor: Subsampled tensor (#batch, time', odim), + where time' = time // 8. + torch.Tensor: Subsampled mask (#batch, 1, time'), + where time' = time // 8. + + """ + x = x.unsqueeze(1) # (b, c, t, f) + x = self.conv(x) + b, c, t, f = x.size() + x = self.out(x.transpose(1, 2).contiguous().view(b, t, c * f)) + if x_mask is None: + return x, None + return x, x_mask[:, :, :-2:2][:, :, :-2:2][:, :, :-2:2] + + +class Conv1dSubsampling(torch.nn.Module): + """Convolutional 1D subsampling (to 1/2 length). + + Args: + idim (int): Input dimension. + odim (int): Output dimension. + dropout_rate (float): Dropout rate. + pos_enc (torch.nn.Module): Custom position encoding layer. + + """ + + def __init__( + self, + idim, + odim, + kernel_size, + stride, + pad, + tf2torch_tensor_name_prefix_torch: str = "stride_conv", + tf2torch_tensor_name_prefix_tf: str = "seq2seq/proj_encoder/downsampling", + ): + super(Conv1dSubsampling, self).__init__() + self.conv = torch.nn.Conv1d(idim, odim, kernel_size, stride) + self.pad_fn = torch.nn.ConstantPad1d(pad, 0.0) + self.stride = stride + self.odim = odim + self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch + self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf + + def output_size(self) -> int: + return self.odim + + def forward(self, x, x_len): + """Subsample x.""" + x = x.transpose(1, 2) # (b, d ,t) + x = self.pad_fn(x) + # x = F.relu(self.conv(x)) + x = F.leaky_relu(self.conv(x), negative_slope=0.0) + x = x.transpose(1, 2) # (b, t ,d) + + if x_len is None: + + return x, None + x_len = (x_len - 1) // self.stride + 1 + return x, x_len + + +class StreamingConvInput(torch.nn.Module): + """Streaming ConvInput module definition. + Args: + input_size: Input size. + conv_size: Convolution size. + subsampling_factor: Subsampling factor. + vgg_like: Whether to use a VGG-like network. + output_size: Block output dimension. + """ + + def __init__( + self, + input_size: int, + conv_size: Union[int, Tuple], + subsampling_factor: int = 4, + vgg_like: bool = True, + conv_kernel_size: int = 3, + output_size: Optional[int] = None, + ) -> None: + """Construct a ConvInput object.""" + super().__init__() + if vgg_like: + if subsampling_factor == 1: + conv_size1, conv_size2 = conv_size + + self.conv = torch.nn.Sequential( + torch.nn.Conv2d( + 1, + conv_size1, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size1, + conv_size1, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.MaxPool2d((1, 2)), + torch.nn.Conv2d( + conv_size1, + conv_size2, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size2, + conv_size2, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.MaxPool2d((1, 2)), + ) + + output_proj = conv_size2 * ((input_size // 2) // 2) + + self.subsampling_factor = 1 + + self.stride_1 = 1 + + self.create_new_mask = self.create_new_vgg_mask + + else: + conv_size1, conv_size2 = conv_size + + kernel_1 = int(subsampling_factor / 2) + + self.conv = torch.nn.Sequential( + torch.nn.Conv2d( + 1, + conv_size1, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size1, + conv_size1, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.MaxPool2d((kernel_1, 2)), + torch.nn.Conv2d( + conv_size1, + conv_size2, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size2, + conv_size2, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.MaxPool2d((2, 2)), + ) + + output_proj = conv_size2 * ((input_size // 2) // 2) + + self.subsampling_factor = subsampling_factor + + self.create_new_mask = self.create_new_vgg_mask + + self.stride_1 = kernel_1 + + else: + if subsampling_factor == 1: + self.conv = torch.nn.Sequential( + torch.nn.Conv2d(1, conv_size, 3, [1, 2], [1, 0]), + torch.nn.ReLU(), + torch.nn.Conv2d(conv_size, conv_size, conv_kernel_size, [1, 2], [1, 0]), + torch.nn.ReLU(), + ) + + output_proj = conv_size * (((input_size - 1) // 2 - 1) // 2) + + self.subsampling_factor = subsampling_factor + self.kernel_2 = conv_kernel_size + self.stride_2 = 1 + + self.create_new_mask = self.create_new_conv2d_mask + + else: + kernel_2, stride_2, conv_2_output_size = sub_factor_to_params( + subsampling_factor, + input_size, + ) + + self.conv = torch.nn.Sequential( + torch.nn.Conv2d(1, conv_size, 3, 2, [1, 0]), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size, conv_size, kernel_2, stride_2, [(kernel_2 - 1) // 2, 0] + ), + torch.nn.ReLU(), + ) + + output_proj = conv_size * conv_2_output_size + + self.subsampling_factor = subsampling_factor + self.kernel_2 = kernel_2 + self.stride_2 = stride_2 + + self.create_new_mask = self.create_new_conv2d_mask + + self.vgg_like = vgg_like + self.min_frame_length = 7 + + if output_size is not None: + self.output = torch.nn.Linear(output_proj, output_size) + self.output_size = output_size + else: + self.output = None + self.output_size = output_proj + + def forward( + self, x: torch.Tensor, mask: Optional[torch.Tensor], chunk_size: Optional[torch.Tensor] + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Encode input sequences. + Args: + x: ConvInput input sequences. (B, T, D_feats) + mask: Mask of input sequences. (B, 1, T) + Returns: + x: ConvInput output sequences. (B, sub(T), D_out) + mask: Mask of output sequences. (B, 1, sub(T)) + """ + if mask is not None: + mask = self.create_new_mask(mask) + olens = max(mask.eq(0).sum(1)) + + b, t, f = x.size() + x = x.unsqueeze(1) # (b. 1. t. f) + + if chunk_size is not None: + max_input_length = int( + chunk_size + * self.subsampling_factor + * (math.ceil(float(t) / (chunk_size * self.subsampling_factor))) + ) + x = map(lambda inputs: pad_to_len(inputs, max_input_length, 1), x) + x = list(x) + x = torch.stack(x, dim=0) + N_chunks = max_input_length // (chunk_size * self.subsampling_factor) + x = x.view(b * N_chunks, 1, chunk_size * self.subsampling_factor, f) + + x = self.conv(x) + + _, c, _, f = x.size() + if chunk_size is not None: + x = x.transpose(1, 2).contiguous().view(b, -1, c * f)[:, :olens, :] + else: + x = x.transpose(1, 2).contiguous().view(b, -1, c * f) + + if self.output is not None: + x = self.output(x) + + return x, mask[:, :olens][:, : x.size(1)] + + def create_new_vgg_mask(self, mask: torch.Tensor) -> torch.Tensor: + """Create a new mask for VGG output sequences. + Args: + mask: Mask of input sequences. (B, T) + Returns: + mask: Mask of output sequences. (B, sub(T)) + """ + if self.subsampling_factor > 1: + vgg1_t_len = mask.size(1) - (mask.size(1) % (self.subsampling_factor // 2)) + mask = mask[:, :vgg1_t_len][:, :: self.subsampling_factor // 2] + + vgg2_t_len = mask.size(1) - (mask.size(1) % 2) + mask = mask[:, :vgg2_t_len][:, ::2] + else: + mask = mask + + return mask + + def create_new_conv2d_mask(self, mask: torch.Tensor) -> torch.Tensor: + """Create new conformer mask for Conv2d output sequences. + Args: + mask: Mask of input sequences. (B, T) + Returns: + mask: Mask of output sequences. (B, sub(T)) + """ + if self.subsampling_factor > 1: + return mask[:, ::2][:, :: self.stride_2] + else: + return mask + + def get_size_before_subsampling(self, size: int) -> int: + """Return the original size before subsampling for a given size. + Args: + size: Number of frames after subsampling. + Returns: + : Number of frames before subsampling. + """ + return size * self.subsampling_factor diff --git a/funasr/models/transformer/utils/subsampling_without_posenc.py b/funasr/models/transformer/utils/subsampling_without_posenc.py new file mode 100644 index 0000000000000000000000000000000000000000..239d3f1ade7f03435e44bf2414f7ab59cb055e6f --- /dev/null +++ b/funasr/models/transformer/utils/subsampling_without_posenc.py @@ -0,0 +1,61 @@ +# Copyright 2020 Emiru Tsunoo +# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +"""Subsampling layer definition.""" + +import math +import torch + + +class Conv2dSubsamplingWOPosEnc(torch.nn.Module): + """Convolutional 2D subsampling. + + Args: + idim (int): Input dimension. + odim (int): Output dimension. + dropout_rate (float): Dropout rate. + kernels (list): kernel sizes + strides (list): stride sizes + + """ + + def __init__(self, idim, odim, dropout_rate, kernels, strides): + """Construct an Conv2dSubsamplingWOPosEnc object.""" + assert len(kernels) == len(strides) + super().__init__() + conv = [] + olen = idim + for i, (k, s) in enumerate(zip(kernels, strides)): + conv += [ + torch.nn.Conv2d(1 if i == 0 else odim, odim, k, s), + torch.nn.ReLU(), + ] + olen = math.floor((olen - k) / s + 1) + self.conv = torch.nn.Sequential(*conv) + self.out = torch.nn.Linear(odim * olen, odim) + self.strides = strides + self.kernels = kernels + + def forward(self, x, x_mask): + """Subsample x. + + Args: + x (torch.Tensor): Input tensor (#batch, time, idim). + x_mask (torch.Tensor): Input mask (#batch, 1, time). + + Returns: + torch.Tensor: Subsampled tensor (#batch, time', odim), + where time' = time // 4. + torch.Tensor: Subsampled mask (#batch, 1, time'), + where time' = time // 4. + + """ + x = x.unsqueeze(1) # (b, c, t, f) + x = self.conv(x) + b, c, t, f = x.size() + x = self.out(x.transpose(1, 2).contiguous().view(b, t, c * f)) + if x_mask is None: + return x, None + for k, s in zip(self.kernels, self.strides): + x_mask = x_mask[:, :, : -k + 1 : s] + return x, x_mask diff --git a/funasr/models/transformer/utils/vgg2l.py b/funasr/models/transformer/utils/vgg2l.py new file mode 100644 index 0000000000000000000000000000000000000000..7466baccf476d6e9395b8a3efd03ec701e9c5a34 --- /dev/null +++ b/funasr/models/transformer/utils/vgg2l.py @@ -0,0 +1,88 @@ +"""VGG2L module definition for custom encoder.""" + +from typing import Tuple, Union + +import torch + + +class VGG2L(torch.nn.Module): + """VGG2L module for custom encoder. + + Args: + idim: Input dimension. + odim: Output dimension. + pos_enc: Positional encoding class. + + """ + + def __init__(self, idim: int, odim: int, pos_enc: torch.nn.Module = None): + """Construct a VGG2L object.""" + super().__init__() + + self.vgg2l = torch.nn.Sequential( + torch.nn.Conv2d(1, 64, 3, stride=1, padding=1), + torch.nn.ReLU(), + torch.nn.Conv2d(64, 64, 3, stride=1, padding=1), + torch.nn.ReLU(), + torch.nn.MaxPool2d((3, 2)), + torch.nn.Conv2d(64, 128, 3, stride=1, padding=1), + torch.nn.ReLU(), + torch.nn.Conv2d(128, 128, 3, stride=1, padding=1), + torch.nn.ReLU(), + torch.nn.MaxPool2d((2, 2)), + ) + + if pos_enc is not None: + self.output = torch.nn.Sequential( + torch.nn.Linear(128 * ((idim // 2) // 2), odim), pos_enc + ) + else: + self.output = torch.nn.Linear(128 * ((idim // 2) // 2), odim) + + def forward(self, feats: torch.Tensor, feats_mask: torch.Tensor) -> Union[ + Tuple[torch.Tensor, torch.Tensor], + Tuple[Tuple[torch.Tensor, torch.Tensor], torch.Tensor], + ]: + """Forward VGG2L bottleneck. + + Args: + feats: Feature sequences. (B, F, D_feats) + feats_mask: Mask of feature sequences. (B, 1, F) + + Returns: + vgg_output: VGG output sequences. + (B, sub(F), D_out) or ((B, sub(F), D_out), (B, sub(F), D_att)) + vgg_mask: Mask of VGG output sequences. (B, 1, sub(F)) + + """ + feats = feats.unsqueeze(1) + vgg_output = self.vgg2l(feats) + + b, c, t, f = vgg_output.size() + + vgg_output = self.output(vgg_output.transpose(1, 2).contiguous().view(b, t, c * f)) + + if feats_mask is not None: + vgg_mask = self.create_new_mask(feats_mask) + else: + vgg_mask = feats_mask + + return vgg_output, vgg_mask + + def create_new_mask(self, feats_mask: torch.Tensor) -> torch.Tensor: + """Create a subsampled mask of feature sequences. + + Args: + feats_mask: Mask of feature sequences. (B, 1, F) + + Returns: + vgg_mask: Mask of VGG2L output sequences. (B, 1, sub(F)) + + """ + vgg1_t_len = feats_mask.size(2) - (feats_mask.size(2) % 3) + vgg_mask = feats_mask[:, :, :vgg1_t_len][:, :, ::3] + + vgg2_t_len = vgg_mask.size(2) - (vgg_mask.size(2) % 2) + vgg_mask = vgg_mask[:, :, :vgg2_t_len][:, :, ::2] + + return vgg_mask diff --git a/funasr/models/uniasr/__init__.py b/funasr/models/uniasr/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/models/uniasr/beam_search.py b/funasr/models/uniasr/beam_search.py new file mode 100644 index 0000000000000000000000000000000000000000..f422129d0883f97736fb640a0fc1330eafe26912 --- /dev/null +++ b/funasr/models/uniasr/beam_search.py @@ -0,0 +1,499 @@ +"""Beam search module.""" + +from itertools import chain +import logging +from typing import Any +from typing import Dict +from typing import List +from typing import NamedTuple +from typing import Tuple +from typing import Union + +import torch + +from funasr.metrics.common import end_detect +from funasr.models.transformer.scorers.scorer_interface import PartialScorerInterface +from funasr.models.transformer.scorers.scorer_interface import ScorerInterface + + +class Hypothesis(NamedTuple): + """Hypothesis data type.""" + + yseq: torch.Tensor + score: Union[float, torch.Tensor] = 0 + scores: Dict[str, Union[float, torch.Tensor]] = dict() + states: Dict[str, Any] = dict() + + def asdict(self) -> dict: + """Convert data to JSON-friendly dict.""" + return self._replace( + yseq=self.yseq.tolist(), + score=float(self.score), + scores={k: float(v) for k, v in self.scores.items()}, + )._asdict() + + +class BeamSearchScama(torch.nn.Module): + """Beam search implementation.""" + + def __init__( + self, + scorers: Dict[str, ScorerInterface], + weights: Dict[str, float], + beam_size: int, + vocab_size: int, + sos: int, + eos: int, + token_list: List[str] = None, + pre_beam_ratio: float = 1.5, + pre_beam_score_key: str = None, + ): + """Initialize beam search. + + Args: + scorers (dict[str, ScorerInterface]): Dict of decoder modules + e.g., Decoder, CTCPrefixScorer, LM + The scorer will be ignored if it is `None` + weights (dict[str, float]): Dict of weights for each scorers + The scorer will be ignored if its weight is 0 + beam_size (int): The number of hypotheses kept during search + vocab_size (int): The number of vocabulary + sos (int): Start of sequence id + eos (int): End of sequence id + token_list (list[str]): List of tokens for debug log + pre_beam_score_key (str): key of scores to perform pre-beam search + pre_beam_ratio (float): beam size in the pre-beam search + will be `int(pre_beam_ratio * beam_size)` + + """ + super().__init__() + # set scorers + self.weights = weights + self.scorers = dict() + self.full_scorers = dict() + self.part_scorers = dict() + # this module dict is required for recursive cast + # `self.to(device, dtype)` in `recog.py` + self.nn_dict = torch.nn.ModuleDict() + for k, v in scorers.items(): + w = weights.get(k, 0) + if w == 0 or v is None: + continue + assert isinstance( + v, ScorerInterface + ), f"{k} ({type(v)}) does not implement ScorerInterface" + self.scorers[k] = v + if isinstance(v, PartialScorerInterface): + self.part_scorers[k] = v + else: + self.full_scorers[k] = v + if isinstance(v, torch.nn.Module): + self.nn_dict[k] = v + + # set configurations + self.sos = sos + self.eos = eos + self.token_list = token_list + self.pre_beam_size = int(pre_beam_ratio * beam_size) + self.beam_size = beam_size + self.n_vocab = vocab_size + if ( + pre_beam_score_key is not None + and pre_beam_score_key != "full" + and pre_beam_score_key not in self.full_scorers + ): + raise KeyError(f"{pre_beam_score_key} is not found in {self.full_scorers}") + self.pre_beam_score_key = pre_beam_score_key + self.do_pre_beam = ( + self.pre_beam_score_key is not None + and self.pre_beam_size < self.n_vocab + and len(self.part_scorers) > 0 + ) + + def init_hyp(self, x: torch.Tensor) -> List[Hypothesis]: + """Get an initial hypothesis data. + + Args: + x (torch.Tensor): The encoder output feature + + Returns: + Hypothesis: The initial hypothesis. + + """ + init_states = dict() + init_scores = dict() + for k, d in self.scorers.items(): + init_states[k] = d.init_state(x) + init_scores[k] = 0.0 + return [ + Hypothesis( + score=0.0, + scores=init_scores, + states=init_states, + yseq=torch.tensor([self.sos], device=x.device), + ) + ] + + @staticmethod + def append_token(xs: torch.Tensor, x: int) -> torch.Tensor: + """Append new token to prefix tokens. + + Args: + xs (torch.Tensor): The prefix token + x (int): The new token to append + + Returns: + torch.Tensor: New tensor contains: xs + [x] with xs.dtype and xs.device + + """ + x = torch.tensor([x], dtype=xs.dtype, device=xs.device) + return torch.cat((xs, x)) + + def score_full( + self, + hyp: Hypothesis, + x: torch.Tensor, + x_mask: torch.Tensor = None, + pre_acoustic_embeds: torch.Tensor = None, + ) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: + """Score new hypothesis by `self.full_scorers`. + + Args: + hyp (Hypothesis): Hypothesis with prefix tokens to score + x (torch.Tensor): Corresponding input feature + + Returns: + Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: Tuple of + score dict of `hyp` that has string keys of `self.full_scorers` + and tensor score values of shape: `(self.n_vocab,)`, + and state dict that has string keys + and state values of `self.full_scorers` + + """ + scores = dict() + states = dict() + for k, d in self.full_scorers.items(): + scores[k], states[k] = d.score( + hyp.yseq, hyp.states[k], x, x_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds + ) + return scores, states + + def score_partial( + self, hyp: Hypothesis, ids: torch.Tensor, x: torch.Tensor + ) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: + """Score new hypothesis by `self.part_scorers`. + + Args: + hyp (Hypothesis): Hypothesis with prefix tokens to score + ids (torch.Tensor): 1D tensor of new partial tokens to score + x (torch.Tensor): Corresponding input feature + + Returns: + Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: Tuple of + score dict of `hyp` that has string keys of `self.part_scorers` + and tensor score values of shape: `(len(ids),)`, + and state dict that has string keys + and state values of `self.part_scorers` + + """ + scores = dict() + states = dict() + for k, d in self.part_scorers.items(): + scores[k], states[k] = d.score_partial(hyp.yseq, ids, hyp.states[k], x) + return scores, states + + def beam( + self, weighted_scores: torch.Tensor, ids: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute topk full token ids and partial token ids. + + Args: + weighted_scores (torch.Tensor): The weighted sum scores for each tokens. + Its shape is `(self.n_vocab,)`. + ids (torch.Tensor): The partial token ids to compute topk + + Returns: + Tuple[torch.Tensor, torch.Tensor]: + The topk full token ids and partial token ids. + Their shapes are `(self.beam_size,)` + + """ + # no pre beam performed + if weighted_scores.size(0) == ids.size(0): + top_ids = weighted_scores.topk(self.beam_size)[1] + return top_ids, top_ids + + # mask pruned in pre-beam not to select in topk + tmp = weighted_scores[ids] + weighted_scores[:] = -float("inf") + weighted_scores[ids] = tmp + top_ids = weighted_scores.topk(self.beam_size)[1] + local_ids = weighted_scores[ids].topk(self.beam_size)[1] + return top_ids, local_ids + + @staticmethod + def merge_scores( + prev_scores: Dict[str, float], + next_full_scores: Dict[str, torch.Tensor], + full_idx: int, + next_part_scores: Dict[str, torch.Tensor], + part_idx: int, + ) -> Dict[str, torch.Tensor]: + """Merge scores for new hypothesis. + + Args: + prev_scores (Dict[str, float]): + The previous hypothesis scores by `self.scorers` + next_full_scores (Dict[str, torch.Tensor]): scores by `self.full_scorers` + full_idx (int): The next token id for `next_full_scores` + next_part_scores (Dict[str, torch.Tensor]): + scores of partial tokens by `self.part_scorers` + part_idx (int): The new token id for `next_part_scores` + + Returns: + Dict[str, torch.Tensor]: The new score dict. + Its keys are names of `self.full_scorers` and `self.part_scorers`. + Its values are scalar tensors by the scorers. + + """ + new_scores = dict() + for k, v in next_full_scores.items(): + new_scores[k] = prev_scores[k] + v[full_idx] + for k, v in next_part_scores.items(): + new_scores[k] = prev_scores[k] + v[part_idx] + return new_scores + + def merge_states(self, states: Any, part_states: Any, part_idx: int) -> Any: + """Merge states for new hypothesis. + + Args: + states: states of `self.full_scorers` + part_states: states of `self.part_scorers` + part_idx (int): The new token id for `part_scores` + + Returns: + Dict[str, torch.Tensor]: The new score dict. + Its keys are names of `self.full_scorers` and `self.part_scorers`. + Its values are states of the scorers. + + """ + new_states = dict() + for k, v in states.items(): + new_states[k] = v + for k, d in self.part_scorers.items(): + new_states[k] = d.select_state(part_states[k], part_idx) + return new_states + + def search( + self, + running_hyps: List[Hypothesis], + x: torch.Tensor, + x_mask: torch.Tensor = None, + pre_acoustic_embeds: torch.Tensor = None, + ) -> List[Hypothesis]: + """Search new tokens for running hypotheses and encoded speech x. + + Args: + running_hyps (List[Hypothesis]): Running hypotheses on beam + x (torch.Tensor): Encoded speech feature (T, D) + + Returns: + List[Hypotheses]: Best sorted hypotheses + + """ + best_hyps = [] + part_ids = torch.arange(self.n_vocab, device=x.device) # no pre-beam + for hyp in running_hyps: + # scoring + weighted_scores = torch.zeros(self.n_vocab, dtype=x.dtype, device=x.device) + scores, states = self.score_full( + hyp, x, x_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds + ) + for k in self.full_scorers: + weighted_scores += self.weights[k] * scores[k] + # partial scoring + if self.do_pre_beam: + pre_beam_scores = ( + weighted_scores + if self.pre_beam_score_key == "full" + else scores[self.pre_beam_score_key] + ) + part_ids = torch.topk(pre_beam_scores, self.pre_beam_size)[1] + part_scores, part_states = self.score_partial(hyp, part_ids, x) + for k in self.part_scorers: + weighted_scores[part_ids] += self.weights[k] * part_scores[k] + # add previous hyp score + weighted_scores += hyp.score + + # update hyps + for j, part_j in zip(*self.beam(weighted_scores, part_ids)): + # will be (2 x beam at most) + best_hyps.append( + Hypothesis( + score=weighted_scores[j], + yseq=self.append_token(hyp.yseq, j), + scores=self.merge_scores(hyp.scores, scores, j, part_scores, part_j), + states=self.merge_states(states, part_states, part_j), + ) + ) + + # sort and prune 2 x beam -> beam + best_hyps = sorted(best_hyps, key=lambda x: x.score, reverse=True)[ + : min(len(best_hyps), self.beam_size) + ] + return best_hyps + + def forward( + self, + x: torch.Tensor, + scama_mask: torch.Tensor = None, + pre_acoustic_embeds: torch.Tensor = None, + maxlenratio: float = 0.0, + minlenratio: float = 0.0, + maxlen: int = None, + minlen: int = 0, + ) -> List[Hypothesis]: + """Perform beam search. + + Args: + x (torch.Tensor): Encoded speech feature (T, D) + maxlenratio (float): Input length ratio to obtain max output length. + If maxlenratio=0.0 (default), it uses a end-detect function + to automatically find maximum hypothesis lengths + If maxlenratio<0.0, its absolute value is interpreted + as a constant max output length. + minlenratio (float): Input length ratio to obtain min output length. + + Returns: + list[Hypothesis]: N-best decoding results + + """ + if maxlen is None: + # set length bounds + if maxlenratio == 0: + maxlen = x.shape[0] + elif maxlenratio < 0: + maxlen = -1 * int(maxlenratio) + else: + maxlen = max(1, int(maxlenratio * x.size(0))) + minlen = int(minlenratio * x.size(0)) + + logging.info("decoder input length: " + str(x.shape[0])) + logging.info("max output length: " + str(maxlen)) + logging.info("min output length: " + str(minlen)) + + # main loop of prefix search + running_hyps = self.init_hyp(x) + ended_hyps = [] + for i in range(maxlen): + logging.debug("position " + str(i)) + mask_enc = None + if scama_mask is not None: + token_num_predictor = scama_mask.size(1) + token_id_slice = min(i, token_num_predictor - 1) + mask_enc = scama_mask[:, token_id_slice : token_id_slice + 1, :] + # if mask_enc.size(1) == 0: + # mask_enc = scama_mask[:, -2:-1, :] + # # mask_enc = torch.zeros_like(mask_enc) + pre_acoustic_embeds_cur = None + if pre_acoustic_embeds is not None: + b, t, d = pre_acoustic_embeds.size() + pad = torch.zeros((b, 1, d), dtype=pre_acoustic_embeds.dtype).to( + device=pre_acoustic_embeds.device + ) + pre_acoustic_embeds = torch.cat((pre_acoustic_embeds, pad), dim=1) + token_id_slice = min(i, t) + pre_acoustic_embeds_cur = pre_acoustic_embeds[ + :, token_id_slice : token_id_slice + 1, : + ] + + best = self.search( + running_hyps, x, x_mask=mask_enc, pre_acoustic_embeds=pre_acoustic_embeds_cur + ) + # post process of one iteration + running_hyps = self.post_process(i, maxlen, maxlenratio, best, ended_hyps) + # end detection + if maxlenratio == 0.0 and end_detect([h.asdict() for h in ended_hyps], i): + logging.info(f"end detected at {i}") + break + if len(running_hyps) == 0: + logging.info("no hypothesis. Finish decoding.") + break + else: + logging.debug(f"remained hypotheses: {len(running_hyps)}") + + nbest_hyps = sorted(ended_hyps, key=lambda x: x.score, reverse=True) + # check the number of hypotheses reaching to eos + if len(nbest_hyps) == 0: + logging.warning( + "there is no N-best results, perform recognition " "again with smaller minlenratio." + ) + return ( + [] + if minlenratio < 0.1 + else self.forward(x, maxlenratio, max(0.0, minlenratio - 0.1)) + ) + + # report the best result + for x in nbest_hyps: + yseq = "".join([self.token_list[x] for x in x.yseq]) + logging.debug("nbest: y: {}, yseq: {}, score: {}".format(x.yseq, yseq, x.score)) + best = nbest_hyps[0] + for k, v in best.scores.items(): + logging.info(f"{v:6.2f} * {self.weights[k]:3} = {v * self.weights[k]:6.2f} for {k}") + logging.info(f"total log probability: {best.score:.2f}") + logging.info(f"normalized log probability: {best.score / len(best.yseq):.2f}") + logging.info(f"total number of ended hypotheses: {len(nbest_hyps)}") + if self.token_list is not None: + logging.info( + "best hypo: " + "".join([self.token_list[x] for x in best.yseq[1:-1]]) + "\n" + ) + return nbest_hyps + + def post_process( + self, + i: int, + maxlen: int, + maxlenratio: float, + running_hyps: List[Hypothesis], + ended_hyps: List[Hypothesis], + ) -> List[Hypothesis]: + """Perform post-processing of beam search iterations. + + Args: + i (int): The length of hypothesis tokens. + maxlen (int): The maximum length of tokens in beam search. + maxlenratio (int): The maximum length ratio in beam search. + running_hyps (List[Hypothesis]): The running hypotheses in beam search. + ended_hyps (List[Hypothesis]): The ended hypotheses in beam search. + + Returns: + List[Hypothesis]: The new running hypotheses. + + """ + logging.debug(f"the number of running hypotheses: {len(running_hyps)}") + if self.token_list is not None: + logging.debug( + "best hypo: " + "".join([self.token_list[x] for x in running_hyps[0].yseq[1:]]) + ) + # add eos in the final loop to avoid that there are no ended hyps + if i == maxlen - 1: + logging.info("adding in the last position in the loop") + running_hyps = [ + h._replace(yseq=self.append_token(h.yseq, self.eos)) for h in running_hyps + ] + + # add ended hypotheses to a final list, and removed them from current hypotheses + # (this will be a problem, number of hyps < beam) + remained_hyps = [] + for hyp in running_hyps: + if hyp.yseq[-1] == self.eos: + # e.g., Word LM needs to add final score + for k, d in chain(self.full_scorers.items(), self.part_scorers.items()): + s = d.final_score(hyp.states[k]) + hyp.scores[k] += s + hyp = hyp._replace(score=hyp.score + self.weights[k] * s) + ended_hyps.append(hyp) + else: + remained_hyps.append(hyp) + return remained_hyps diff --git a/funasr/models/uniasr/model.py b/funasr/models/uniasr/model.py new file mode 100644 index 0000000000000000000000000000000000000000..bde637777ef27d694dd2f19260828ba2d60da317 --- /dev/null +++ b/funasr/models/uniasr/model.py @@ -0,0 +1,1030 @@ +#!/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 + 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/funasr/models/uniasr/template.yaml b/funasr/models/uniasr/template.yaml new file mode 100644 index 0000000000000000000000000000000000000000..43d55fc265b0f3aba301308283786c46cca73a25 --- /dev/null +++ b/funasr/models/uniasr/template.yaml @@ -0,0 +1,206 @@ +# This is an example that demonstrates how to configure a model file. +# You can modify the configuration according to your own requirements. + +# to print the register_table: +# from funasr.register import tables +# tables.print() + +# network architecture +model: UniASR +model_conf: + ctc_weight: 0.0 + lsm_weight: 0.1 + length_normalized_loss: true + predictor_weight: 1.0 + decoder_attention_chunk_type: chunk + ctc_weight2: 0.0 + predictor_weight2: 1.0 + decoder_attention_chunk_type2: chunk + loss_weight_model1: 0.5 + + +# encoder +encoder: SANMEncoderChunkOpt +encoder_conf: + output_size: 320 + attention_heads: 4 + linear_units: 1280 + num_blocks: 35 + dropout_rate: 0.1 + positional_dropout_rate: 0.1 + attention_dropout_rate: 0.1 + input_layer: pe + pos_enc_class: SinusoidalPositionEncoder + normalize_before: true + kernel_size: 11 + sanm_shfit: 0 + selfattention_layer_type: sanm + chunk_size: + - 20 + - 60 + stride: + - 10 + - 40 + pad_left: + - 5 + - 10 + encoder_att_look_back_factor: + - 0 + - 0 + decoder_att_look_back_factor: + - 0 + - 0 + +# decoder +decoder: FsmnDecoderSCAMAOpt +decoder_conf: + attention_dim: 256 + attention_heads: 4 + linear_units: 1024 + num_blocks: 12 + dropout_rate: 0.1 + positional_dropout_rate: 0.1 + self_attention_dropout_rate: 0.1 + src_attention_dropout_rate: 0.1 + att_layer_num: 6 + kernel_size: 11 + concat_embeds: true + +# predictor +predictor: CifPredictorV2 +predictor_conf: + idim: 320 + threshold: 1.0 + l_order: 1 + r_order: 1 + + +# encoder2 +encoder2: SANMEncoderChunkOpt +encoder2_conf: + output_size: 320 + attention_heads: 4 + linear_units: 1280 + num_blocks: 20 + dropout_rate: 0.1 + positional_dropout_rate: 0.1 + attention_dropout_rate: 0.1 + input_layer: pe + pos_enc_class: SinusoidalPositionEncoder + normalize_before: true + kernel_size: 21 + sanm_shfit: 0 + selfattention_layer_type: sanm + chunk_size: + - 45 + - 70 + stride: + - 35 + - 50 + pad_left: + - 5 + - 10 + encoder_att_look_back_factor: + - 0 + - 0 + decoder_att_look_back_factor: + - 0 + - 0 + +# decoder +decoder2: FsmnDecoderSCAMAOpt +decoder2_conf: + attention_dim: 320 + attention_heads: 4 + linear_units: 1280 + num_blocks: 12 + dropout_rate: 0.1 + positional_dropout_rate: 0.1 + self_attention_dropout_rate: 0.1 + src_attention_dropout_rate: 0.1 + att_layer_num: 6 + kernel_size: 11 + concat_embeds: true + +predictor2: CifPredictorV2 +predictor2_conf: + idim: 320 + threshold: 1.0 + l_order: 1 + r_order: 1 + +stride_conv: stride_conv1d +stride_conv_conf: + kernel_size: 2 + stride: 2 + pad: + - 0 + - 1 + +# frontend related +frontend: WavFrontend +frontend_conf: + fs: 16000 + window: hamming + n_mels: 80 + frame_length: 25 + frame_shift: 10 + lfr_m: 7 + lfr_n: 6 + dither: 0.0 + +specaug: SpecAugLFR +specaug_conf: + apply_time_warp: false + time_warp_window: 5 + time_warp_mode: bicubic + apply_freq_mask: true + freq_mask_width_range: + - 0 + - 30 + lfr_rate: 6 + num_freq_mask: 1 + apply_time_mask: true + time_mask_width_range: + - 0 + - 12 + num_time_mask: 1 + +train_conf: + accum_grad: 1 + grad_clip: 5 + max_epoch: 150 + keep_nbest_models: 10 + avg_nbest_model: 10 + log_interval: 50 + +optim: adam +optim_conf: + lr: 0.0001 +scheduler: warmuplr +scheduler_conf: + warmup_steps: 30000 + +dataset: AudioDataset +dataset_conf: + index_ds: IndexDSJsonl + batch_sampler: BatchSampler + batch_type: example # example or length + batch_size: 1 # if batch_type is example, batch_size is the numbers of samples; if length, batch_size is source_token_len+target_token_len; + max_token_length: 2048 # filter samples if source_token_len+target_token_len > max_token_length, + buffer_size: 500 + shuffle: True + num_workers: 0 + +tokenizer: CharTokenizer +tokenizer_conf: + unk_symbol: + split_with_space: true + + +ctc_conf: + dropout_rate: 0.0 + ctc_type: builtin + reduce: true + ignore_nan_grad: true +normalize: null diff --git a/funasr/models/whisper/__init__.py b/funasr/models/whisper/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/models/whisper/model.py b/funasr/models/whisper/model.py new file mode 100644 index 0000000000000000000000000000000000000000..398eea3fda93ebef0267aeef4cab94a5efb769f7 --- /dev/null +++ b/funasr/models/whisper/model.py @@ -0,0 +1,125 @@ +from dataclasses import dataclass +from typing import Dict +from typing import Iterable, Optional +import time +import numpy as np +import torch +import torch.nn.functional as F +from torch import Tensor +from torch import nn + +import whisper + +# import whisper_timestamped as whisper + +from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank + +from funasr.register import tables + + +@tables.register("model_classes", "Whisper-tiny.en") +@tables.register("model_classes", "Whisper-tiny") +@tables.register("model_classes", "Whisper-base.en") +@tables.register("model_classes", "Whisper-base") +@tables.register("model_classes", "Whisper-small.en") +@tables.register("model_classes", "Whisper-small") +@tables.register("model_classes", "Whisper-medium.en") +@tables.register("model_classes", "Whisper-medium") +@tables.register("model_classes", "Whisper-large-v1") +@tables.register("model_classes", "Whisper-large-v2") +@tables.register("model_classes", "Whisper-large-v3") +@tables.register("model_classes", "Whisper-large-v3-turbo") +@tables.register("model_classes", "WhisperWarp") +class WhisperWarp(nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + hub = kwargs.get("hub", "funasr") + if hub == "openai": + model_or_path = kwargs.get("model_path", "Whisper-large-v3") + if model_or_path.startswith("Whisper-"): + model_or_path = model_or_path.replace("Whisper-", "") + model = whisper.load_model(model_or_path) + else: + dims = kwargs.get("dims", {}) + dims = whisper.model.ModelDimensions(**dims) + model = whisper.model.Whisper(dims=dims) + + self.model = model + + self.encoder_output_size = self.model.dims.n_audio_state + + def forward( + self, + ): + pass + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): + if kwargs.get("batch_size", 1) > 1: + raise NotImplementedError("batch decoding is not implemented") + + if frontend is None and not hasattr(self, "frontend"): + frontend_class = tables.frontend_classes.get("WhisperFrontend") + frontend = frontend_class( + n_mels=self.model.dims.n_mels, do_pad_trim=kwargs.get("do_pad_trim", True) + ) + self.frontend = frontend + else: + frontend = frontend if frontend is not None else self.frontend + + 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 if hasattr(frontend, "fs") else 16000, + 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}" + frame_shift = frontend.frame_shift if hasattr(frontend, "frame_shift") else 10 + lfr_n = frontend.lfr_n if hasattr(frontend, "lfr_n") else 1 + meta_data["batch_data_time"] = speech_lengths.sum().item() * frame_shift * lfr_n / 1000 + + speech = speech.to(device=kwargs["device"])[0, :, :] + speech_lengths = speech_lengths.to(device=kwargs["device"]) + + # # detect the spoken language + # _, probs = self.model.detect_language(speech) + # print(f"Detected language: {max(probs, key=probs.get)}") + + # decode the audio + options = whisper.DecodingOptions(**kwargs.get("DecodingOptions", {})) + + result = whisper.decode(self.model, speech, options=options) + # result = whisper.transcribe(self.model, speech) + + results = [] + result_i = {"key": key[0], "text": result.text} + + results.append(result_i) + + return results, meta_data diff --git a/funasr/models/whisper/template.yaml b/funasr/models/whisper/template.yaml new file mode 100644 index 0000000000000000000000000000000000000000..568b9484f121b9fa07b06ea4b74160fe23c9da29 --- /dev/null +++ b/funasr/models/whisper/template.yaml @@ -0,0 +1,45 @@ +# This is an example that demonstrates how to configure a model file. +# You can modify the configuration according to your own requirements. + +# to print the register_table: +# from funasr.register import tables +# tables.print() + +# network architecture +model: WhisperWarp +model_conf: + lsm_weight: 0.1 + length_normalized_loss: true + hub: funasr # openai + + + +# only use for hub == funasr, +# if hub == openai, dims is automaticall download +dims: + n_mels: 128 + n_vocab: 51866 + n_audio_ctx: 1500 + n_audio_state: 1280 + n_audio_head: 20 + n_audio_layer: 32 + n_text_ctx: 448 + n_text_state: 1280 + n_text_head: 20 + n_text_layer: 32 + +# frontend related +frontend: WhisperFrontend +frontend_conf: + fs: 16000 + n_mels: ${dims.n_mels} + do_pad_trim: true + +tokenizer: WhisperTokenizer +tokenizer_conf: + language: null + task: transcribe + is_multilingual: true + num_languages: 100 + +scope_map: [none, "model."] \ No newline at end of file diff --git a/funasr/models/whisper_lid/__init__.py b/funasr/models/whisper_lid/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/models/whisper_lid/decoder.py b/funasr/models/whisper_lid/decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..1c8ab47f862e5cb48ea295332d890da1ea15b747 --- /dev/null +++ b/funasr/models/whisper_lid/decoder.py @@ -0,0 +1,162 @@ +# Copyright ESPnet (https://github.com/espnet/espnet). All Rights Reserved. +# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +import copy +from typing import Any, List, Tuple + +import torch +from torch import nn +import whisper + +from funasr.models.transformer.utils.nets_utils import make_pad_mask +from funasr.register import tables + + +@tables.register("decoder_classes", "OpenAIWhisperDecoderWarp") +class OpenAIWhisperDecoderWarp(nn.Module): + """Transformer-based Speech-to-Text Decoder from OpenAI's Whisper Model: + + URL: https://github.com/openai/whisper + """ + + def __init__( + self, + dropout_rate: float = 0.0, + whisper_model: str = "small", + download_dir: str = None, + use_padmask: bool = False, + ): + super().__init__() + + assert whisper_model in whisper.available_models() + _model = whisper.load_model(whisper_model, download_root=download_dir, device="cpu") + self.decoders = copy.deepcopy(_model.decoder) + attention_dim = self.decoders.token_embedding.embedding_dim + + # note that originally Whisper doesn't use dropouts + self.dropout = torch.nn.Dropout(dropout_rate) + + self.decoders.train() + del _model + self.use_padmask = use_padmask + + def forward( + self, + hs_pad: torch.Tensor, + hlens: torch.Tensor, + ys_in_pad: torch.Tensor, + ys_in_lens: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Forward decoder. + + Args: + hs_pad: encoded memory, float32 (batch, maxlen_in, feat) + hlens: (batch) + ys_in_pad: + input token ids, int64 (batch, maxlen_out) + if input_layer == "embed" + input tensor (batch, maxlen_out, #mels) in the other cases + ys_in_lens: (batch) + Returns: + (tuple): tuple containing: + + x: decoded token score before softmax (batch, maxlen_out, token) + if use_output_layer is True, + olens: (batch, ) + """ + tgt, memory = ys_in_pad, hs_pad + tgt = self.decoders.token_embedding(tgt) + self.decoders.positional_embedding[: tgt.size(1)] + tgt = self.dropout(tgt) + + x = tgt.to(memory.dtype) + + if self.use_padmask: + memory_mask = (~make_pad_mask(hlens)[:, None, :]).to(memory.device) + else: + memory_mask = None + + for layer, block in enumerate(self.decoders.blocks): + x = block( + x, + memory, + mask=self.decoders.mask, + memory_mask=memory_mask, + is_pad_mask=False, + is_pad_memory_mask=True, + ) + + if layer < len(self.decoders.blocks) - 1: + x = self.dropout(x) + + x = self.decoders.ln(x) + x = (x @ torch.transpose(self.decoders.token_embedding.weight.to(x.dtype), 0, 1)).float() + + return x, ys_in_lens + + def forward_one_step( + self, + tgt: torch.Tensor, + tgt_mask: torch.Tensor, + memory: torch.Tensor, + cache: List[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, List[torch.Tensor]]: + """Forward one step. + + Args: + tgt: input token ids, int64 (batch, maxlen_out) + tgt_mask: input token mask, (batch, maxlen_out) + dtype=torch.uint8 in PyTorch 1.2- + dtype=torch.bool in PyTorch 1.2+ (include 1.2) + memory: encoded memory, float32 (batch, maxlen_in, feat) + cache: cached output list of (batch, max_time_out-1, size) + Returns: + y, cache: NN output value and cache per `self.decoders`. + y.shape` is (batch, maxlen_out, token) + NOTE (Shih-Lun): + cache implementation is ignored for now + for simplicity & correctness + """ + x = self.decoders.token_embedding(tgt) + self.decoders.positional_embedding[: tgt.size(1)] + x = self.dropout(x) + x = x.to(memory.dtype) + + for layer, block in enumerate(self.decoders.blocks): + x = block(x, memory, mask=self.decoders.mask) + if layer < len(self.decoders.blocks) - 1: + x = self.dropout(x) + + x = self.decoders.ln(x) + y = x[:, -1] + y = (y @ torch.transpose(self.decoders.token_embedding.weight.to(x.dtype), 0, 1)).float() + y = torch.log_softmax(y, dim=-1) + + return y, None + + def score(self, ys, state, x): + """Score.""" + logp, state = self.forward_one_step( + ys.unsqueeze(0), torch.empty(0), x.unsqueeze(0), cache=state # dummy mask + ) + return logp.squeeze(0), state + + def batch_score( + self, ys: torch.Tensor, states: List[Any], xs: torch.Tensor + ) -> Tuple[torch.Tensor, List[Any]]: + """Score new token batch. + + Args: + ys (torch.Tensor): torch.int64 prefix tokens (n_batch, ylen). + states (List[Any]): Scorer states for prefix tokens. + xs (torch.Tensor): + The encoder feature that generates ys (n_batch, xlen, n_feat). + + Returns: + tuple[torch.Tensor, List[Any]]: Tuple of + batchfied scores for next token with shape of `(n_batch, n_vocab)` + and next state list for ys. + + """ + # batch decoding, dummy mask is passed + logp, states = self.forward_one_step(ys, torch.empty(0), xs, cache=None) + + return logp, None diff --git a/funasr/models/whisper_lid/encoder.py b/funasr/models/whisper_lid/encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..b6fb28b50df6c9992ca6393238e3effeb7aa541a --- /dev/null +++ b/funasr/models/whisper_lid/encoder.py @@ -0,0 +1,113 @@ +# Copyright ESPnet (https://github.com/espnet/espnet). All Rights Reserved. +# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +import copy +from typing import Optional, Tuple, Union + +import torch +from torch import nn +import torch.nn.functional as F +import whisper + +from funasr.models.transformer.utils.nets_utils import make_pad_mask +from funasr.models.specaug.specaug import SpecAug +from funasr.register import tables + + +@tables.register("encoder_classes", "OpenAIWhisperEncoderWarp") +class OpenAIWhisperEncoderWarp(nn.Module): + """Transformer-based Speech Encoder from OpenAI's Whisper Model: + + URL: https://github.com/openai/whisper + """ + + def __init__( + self, + dropout_rate: float = 0.0, + whisper_model: str = "small", + download_dir: str = None, + use_specaug: bool = False, + use_padmask: bool = False, + specaug_conf: Union[dict, None] = None, + ): + super().__init__() + + # note that originally Whisper doesn't use dropouts + self.dropout = torch.nn.Dropout(dropout_rate) + + assert whisper_model in whisper.available_models() + _model = whisper.load_model(whisper_model, download_root=download_dir, device="cpu") + self.encoders = copy.deepcopy(_model.encoder) + self.encoders.train() + + del _model + + if use_specaug: + self.specaug = SpecAug(**specaug_conf) + else: + self.specaug = None + self.use_padmask = use_padmask + + def whisper_encode( + self, + input: torch.Tensor, + ilens: torch.Tensor = None, + ) -> torch.Tensor: + x = F.gelu(self.encoders.conv1(input)) + x = F.gelu(self.encoders.conv2(x)) + x = x.permute(0, 2, 1) + + n_frames = x.size(1) + max_pos = self.encoders.positional_embedding.size(0) + if n_frames <= max_pos: + x = (x + self.encoders.positional_embedding[: x.size(1), :]).to(x.dtype) + else: + # due to positional encoding, audios >30 sec won't be accepted + x = x[:, :max_pos, :] + self.encoders.positional_embedding + + if ilens is not None: + olens = ( + 1 + + (ilens - self.encoders.conv2.kernel_size[0] + 2 * self.encoders.conv2.padding[0]) + // self.encoders.conv2.stride[0] + ) + olens = torch.clamp(olens, max=max_pos) + else: + olens = None + + if self.use_padmask: + padding_mask = (~make_pad_mask(olens)[:, None, :]).to(x.device) + else: + padding_mask = None + + x = self.dropout(x) + + for layer, block in enumerate(self.encoders.blocks): + x = block(x) + if layer < len(self.encoders.blocks) - 1: + x = self.dropout(x) + + x = self.encoders.ln_post(x) + + return x, olens + + def output_size(self) -> int: + # dummy output size + return self.encoders.conv2.weight.shape[0] + + def forward( + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + prev_states: torch.Tensor = None, + ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + feats, feats_lens = xs_pad, ilens + + if self.specaug is not None and self.encoders.training: + feats = torch.transpose(feats, 1, 2) + feats, feats_lens = self.specaug(feats, feats_lens) + feats = torch.transpose(feats, 1, 2) + + xs_pad, olens = self.whisper_encode(feats, feats_lens) + + return xs_pad, olens, None diff --git a/funasr/models/whisper_lid/eres2net/ResNet.py b/funasr/models/whisper_lid/eres2net/ResNet.py new file mode 100644 index 0000000000000000000000000000000000000000..df3b96d242eb901890645a8310d146c606d28771 --- /dev/null +++ b/funasr/models/whisper_lid/eres2net/ResNet.py @@ -0,0 +1,387 @@ +# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +""" Res2Net implementation is adapted from https://github.com/wenet-e2e/wespeaker. + ERes2Net incorporates both local and global feature fusion techniques to improve the performance. + The local feature fusion (LFF) fuses the features within one single residual block to extract the local signal. + The global feature fusion (GFF) takes acoustic features of different scales as input to aggregate global signal. + ERes2Net-Large is an upgraded version of ERes2Net that uses a larger number of parameters to achieve better + recognition performance. Parameters expansion, baseWidth, and scale can be modified to obtain optimal performance. +""" + +import torch +import math +import torch.nn as nn +import torch.nn.functional as F +import funasr.models.whisper_lid.eres2net.pooling_layers as pooling_layers +from funasr.models.whisper_lid.eres2net.fusion import AFF + + +class ReLU(nn.Hardtanh): + + def __init__(self, inplace=False): + super(ReLU, self).__init__(0, 20, inplace) + + def __repr__(self): + inplace_str = "inplace" if self.inplace else "" + return self.__class__.__name__ + " (" + inplace_str + ")" + + +def conv1x1(in_planes, out_planes, stride=1): + "1x1 convolution without padding" + return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, padding=0, bias=False) + + +def conv3x3(in_planes, out_planes, stride=1): + "3x3 convolution with padding" + return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False) + + +class BasicBlockERes2Net(nn.Module): + expansion = 2 + + def __init__(self, in_planes, planes, stride=1, baseWidth=32, scale=2): + super(BasicBlockERes2Net, self).__init__() + width = int(math.floor(planes * (baseWidth / 64.0))) + self.conv1 = conv1x1(in_planes, width * scale, stride) + self.bn1 = nn.BatchNorm2d(width * scale) + self.nums = scale + + convs = [] + bns = [] + for i in range(self.nums): + convs.append(conv3x3(width, width)) + bns.append(nn.BatchNorm2d(width)) + self.convs = nn.ModuleList(convs) + self.bns = nn.ModuleList(bns) + self.relu = ReLU(inplace=True) + + self.conv3 = conv1x1(width * scale, planes * self.expansion) + self.bn3 = nn.BatchNorm2d(planes * self.expansion) + self.shortcut = nn.Sequential() + if stride != 1 or in_planes != self.expansion * planes: + self.shortcut = nn.Sequential( + nn.Conv2d( + in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False + ), + nn.BatchNorm2d(self.expansion * planes), + ) + self.stride = stride + self.width = width + self.scale = scale + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + spx = torch.split(out, self.width, 1) + for i in range(self.nums): + if i == 0: + sp = spx[i] + else: + sp = sp + spx[i] + sp = self.convs[i](sp) + sp = self.relu(self.bns[i](sp)) + if i == 0: + out = sp + else: + out = torch.cat((out, sp), 1) + + out = self.conv3(out) + out = self.bn3(out) + + residual = self.shortcut(x) + out += residual + out = self.relu(out) + + return out + + +class BasicBlockERes2Net_diff_AFF(nn.Module): + expansion = 2 + + def __init__(self, in_planes, planes, stride=1, baseWidth=32, scale=2): + super(BasicBlockERes2Net_diff_AFF, self).__init__() + width = int(math.floor(planes * (baseWidth / 64.0))) + self.conv1 = conv1x1(in_planes, width * scale, stride) + self.bn1 = nn.BatchNorm2d(width * scale) + self.nums = scale + + convs = [] + fuse_models = [] + bns = [] + for i in range(self.nums): + convs.append(conv3x3(width, width)) + bns.append(nn.BatchNorm2d(width)) + for j in range(self.nums - 1): + fuse_models.append(AFF(channels=width)) + + self.convs = nn.ModuleList(convs) + self.bns = nn.ModuleList(bns) + self.fuse_models = nn.ModuleList(fuse_models) + self.relu = ReLU(inplace=True) + + self.conv3 = conv1x1(width * scale, planes * self.expansion) + self.bn3 = nn.BatchNorm2d(planes * self.expansion) + self.shortcut = nn.Sequential() + if stride != 1 or in_planes != self.expansion * planes: + self.shortcut = nn.Sequential( + nn.Conv2d( + in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False + ), + nn.BatchNorm2d(self.expansion * planes), + ) + self.stride = stride + self.width = width + self.scale = scale + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + spx = torch.split(out, self.width, 1) + for i in range(self.nums): + if i == 0: + sp = spx[i] + else: + sp = self.fuse_models[i - 1](sp, spx[i]) + + sp = self.convs[i](sp) + sp = self.relu(self.bns[i](sp)) + if i == 0: + out = sp + else: + out = torch.cat((out, sp), 1) + + out = self.conv3(out) + out = self.bn3(out) + + residual = self.shortcut(x) + out += residual + out = self.relu(out) + + return out + + +class ERes2Net(nn.Module): + def __init__( + self, + block=BasicBlockERes2Net, + block_fuse=BasicBlockERes2Net_diff_AFF, + num_blocks=[3, 4, 6, 3], + m_channels=32, + feat_dim=80, + embedding_size=192, + pooling_func="TSTP", + two_emb_layer=False, + ): + super(ERes2Net, self).__init__() + self.in_planes = m_channels + self.feat_dim = feat_dim + self.embedding_size = embedding_size + self.stats_dim = int(feat_dim / 8) * m_channels * 8 + self.two_emb_layer = two_emb_layer + self._output_size = embedding_size + + self.conv1 = nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False) + self.bn1 = nn.BatchNorm2d(m_channels) + self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=1) + self.layer2 = self._make_layer(block, m_channels * 2, num_blocks[1], stride=2) + self.layer3 = self._make_layer(block_fuse, m_channels * 4, num_blocks[2], stride=2) + self.layer4 = self._make_layer(block_fuse, m_channels * 8, num_blocks[3], stride=2) + + # Downsampling module for each layer + self.layer1_downsample = nn.Conv2d( + m_channels * 2, m_channels * 4, kernel_size=3, stride=2, padding=1, bias=False + ) + self.layer2_downsample = nn.Conv2d( + m_channels * 4, m_channels * 8, kernel_size=3, padding=1, stride=2, bias=False + ) + self.layer3_downsample = nn.Conv2d( + m_channels * 8, m_channels * 16, kernel_size=3, padding=1, stride=2, bias=False + ) + + # Bottom-up fusion module + self.fuse_mode12 = AFF(channels=m_channels * 4) + self.fuse_mode123 = AFF(channels=m_channels * 8) + self.fuse_mode1234 = AFF(channels=m_channels * 16) + + self.n_stats = 1 if pooling_func == "TAP" or pooling_func == "TSDP" else 2 + self.pool = getattr(pooling_layers, pooling_func)(in_dim=self.stats_dim * block.expansion) + self.seg_1 = nn.Linear(self.stats_dim * block.expansion * self.n_stats, embedding_size) + if self.two_emb_layer: + self.seg_bn_1 = nn.BatchNorm1d(embedding_size, affine=False) + self.seg_2 = nn.Linear(embedding_size, embedding_size) + else: + self.seg_bn_1 = nn.Identity() + self.seg_2 = nn.Identity() + + def _make_layer(self, block, planes, num_blocks, stride): + strides = [stride] + [1] * (num_blocks - 1) + layers = [] + for stride in strides: + layers.append(block(self.in_planes, planes, stride)) + self.in_planes = planes * block.expansion + return nn.Sequential(*layers) + + def output_size(self) -> int: + return self._output_size + + def forward(self, x, ilens): + # assert x.shape[1] == ilens.max() + x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T) + x = x.unsqueeze_(1) + out = F.relu(self.bn1(self.conv1(x))) + out1 = self.layer1(out) + out2 = self.layer2(out1) + out1_downsample = self.layer1_downsample(out1) + fuse_out12 = self.fuse_mode12(out2, out1_downsample) + out3 = self.layer3(out2) + fuse_out12_downsample = self.layer2_downsample(fuse_out12) + fuse_out123 = self.fuse_mode123(out3, fuse_out12_downsample) + out4 = self.layer4(out3) + fuse_out123_downsample = self.layer3_downsample(fuse_out123) + fuse_out1234 = self.fuse_mode1234(out4, fuse_out123_downsample) + olens = (((((ilens - 1) // 2 + 1) - 1) // 2 + 1) - 1) // 2 + 1 + stats = self.pool(fuse_out1234, olens) + + embed_a = self.seg_1(stats) + if self.two_emb_layer: + out = F.relu(embed_a) + out = self.seg_bn_1(out) + embed_b = self.seg_2(out) + return embed_b + else: + return embed_a + + +class BasicBlockRes2Net(nn.Module): + expansion = 2 + + def __init__(self, in_planes, planes, stride=1, baseWidth=32, scale=2): + super(BasicBlockRes2Net, self).__init__() + width = int(math.floor(planes * (baseWidth / 64.0))) + self.conv1 = conv1x1(in_planes, width * scale, stride) + self.bn1 = nn.BatchNorm2d(width * scale) + self.nums = scale - 1 + convs = [] + bns = [] + for i in range(self.nums): + convs.append(conv3x3(width, width)) + bns.append(nn.BatchNorm2d(width)) + self.convs = nn.ModuleList(convs) + self.bns = nn.ModuleList(bns) + self.relu = ReLU(inplace=True) + + self.conv3 = conv1x1(width * scale, planes * self.expansion) + self.bn3 = nn.BatchNorm2d(planes * self.expansion) + self.shortcut = nn.Sequential() + if stride != 1 or in_planes != self.expansion * planes: + self.shortcut = nn.Sequential( + nn.Conv2d( + in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False + ), + nn.BatchNorm2d(self.expansion * planes), + ) + self.stride = stride + self.width = width + self.scale = scale + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + spx = torch.split(out, self.width, 1) + for i in range(self.nums): + if i == 0: + sp = spx[i] + else: + sp = sp + spx[i] + sp = self.convs[i](sp) + sp = self.relu(self.bns[i](sp)) + if i == 0: + out = sp + else: + out = torch.cat((out, sp), 1) + + out = torch.cat((out, spx[self.nums]), 1) + + out = self.conv3(out) + out = self.bn3(out) + + residual = self.shortcut(x) + out += residual + out = self.relu(out) + + return out + + +class Res2Net(nn.Module): + def __init__( + self, + block=BasicBlockRes2Net, + num_blocks=[3, 4, 6, 3], + m_channels=32, + feat_dim=80, + embedding_size=192, + pooling_func="TSTP", + two_emb_layer=False, + ): + super(Res2Net, self).__init__() + self.in_planes = m_channels + self.feat_dim = feat_dim + self.embedding_size = embedding_size + self.stats_dim = int(feat_dim / 8) * m_channels * 8 + self.two_emb_layer = two_emb_layer + + self.conv1 = nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False) + self.bn1 = nn.BatchNorm2d(m_channels) + self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=1) + self.layer2 = self._make_layer(block, m_channels * 2, num_blocks[1], stride=2) + self.layer3 = self._make_layer(block, m_channels * 4, num_blocks[2], stride=2) + self.layer4 = self._make_layer(block, m_channels * 8, num_blocks[3], stride=2) + + self.n_stats = 1 if pooling_func == "TAP" or pooling_func == "TSDP" else 2 + self.pool = getattr(pooling_layers, pooling_func)(in_dim=self.stats_dim * block.expansion) + self.seg_1 = nn.Linear(self.stats_dim * block.expansion * self.n_stats, embedding_size) + if self.two_emb_layer: + self.seg_bn_1 = nn.BatchNorm1d(embedding_size, affine=False) + self.seg_2 = nn.Linear(embedding_size, embedding_size) + else: + self.seg_bn_1 = nn.Identity() + self.seg_2 = nn.Identity() + + def _make_layer(self, block, planes, num_blocks, stride): + strides = [stride] + [1] * (num_blocks - 1) + layers = [] + for stride in strides: + layers.append(block(self.in_planes, planes, stride)) + self.in_planes = planes * block.expansion + return nn.Sequential(*layers) + + def forward(self, x): + x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T) + + x = x.unsqueeze_(1) + out = F.relu(self.bn1(self.conv1(x))) + out = self.layer1(out) + out = self.layer2(out) + out = self.layer3(out) + out = self.layer4(out) + + stats = self.pool(out) + + embed_a = self.seg_1(stats) + if self.two_emb_layer: + out = F.relu(embed_a) + out = self.seg_bn_1(out) + embed_b = self.seg_2(out) + return embed_b + else: + return embed_a diff --git a/funasr/models/whisper_lid/eres2net/__init__.py b/funasr/models/whisper_lid/eres2net/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/models/whisper_lid/eres2net/fusion.py b/funasr/models/whisper_lid/eres2net/fusion.py new file mode 100644 index 0000000000000000000000000000000000000000..fbe699efbaa33380594df2c5068b78e0ecbe7303 --- /dev/null +++ b/funasr/models/whisper_lid/eres2net/fusion.py @@ -0,0 +1,28 @@ +# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +import torch +import torch.nn as nn + + +class AFF(nn.Module): + + def __init__(self, channels=64, r=4): + super(AFF, self).__init__() + inter_channels = int(channels // r) + + self.local_att = nn.Sequential( + nn.Conv2d(channels * 2, inter_channels, kernel_size=1, stride=1, padding=0), + nn.BatchNorm2d(inter_channels), + nn.SiLU(inplace=True), + nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0), + nn.BatchNorm2d(channels), + ) + + def forward(self, x, ds_y): + xa = torch.cat((x, ds_y), dim=1) + x_att = self.local_att(xa) + x_att = 1.0 + torch.tanh(x_att) + xo = torch.mul(x, x_att) + torch.mul(ds_y, 2.0 - x_att) + + return xo diff --git a/funasr/models/whisper_lid/eres2net/pooling_layers.py b/funasr/models/whisper_lid/eres2net/pooling_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..da23d7b4caf65164aadaae13ad95f8dd7adb288a --- /dev/null +++ b/funasr/models/whisper_lid/eres2net/pooling_layers.py @@ -0,0 +1,119 @@ +# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +""" This implementation is adapted from https://github.com/wenet-e2e/wespeaker.""" + +import torch +import torch.nn as nn + +from funasr.models.transformer.utils.nets_utils import make_pad_mask + + +class TAP(nn.Module): + """ + Temporal average pooling, only first-order mean is considered + """ + + def __init__(self, **kwargs): + super(TAP, self).__init__() + + def forward(self, x): + pooling_mean = x.mean(dim=-1) + # To be compatable with 2D input + pooling_mean = pooling_mean.flatten(start_dim=1) + return pooling_mean + + +class TSDP(nn.Module): + """ + Temporal standard deviation pooling, only second-order std is considered + """ + + def __init__(self, **kwargs): + super(TSDP, self).__init__() + + def forward(self, x): + # The last dimension is the temporal axis + pooling_std = torch.sqrt(torch.var(x, dim=-1) + 1e-8) + pooling_std = pooling_std.flatten(start_dim=1) + return pooling_std + + +class TSTP(nn.Module): + """ + Temporal statistics pooling, concatenate mean and std, which is used in + x-vector + Comment: simple concatenation can not make full use of both statistics + """ + + def __init__(self, **kwargs): + super(TSTP, self).__init__() + + def forward(self, x, olens): + # The last dimension is the temporal axis + masks = (~make_pad_mask(olens, maxlen=x.shape[-1])[:, None, None, :]).to(x.device) + x_masked = x * masks + sum_without_padding = torch.sum(x_masked, axis=-1) + count_without_padding = torch.sum(masks, axis=-1) + mean_without_padding = sum_without_padding / count_without_padding + + var_without_padding = ((x_masked - mean_without_padding.unsqueeze(-1)) ** 2 * masks).sum( + -1 + ) / count_without_padding + + pooling_mean = mean_without_padding + pooling_std = torch.sqrt(var_without_padding + 1e-8) + pooling_mean = pooling_mean.flatten(start_dim=1) + pooling_std = pooling_std.flatten(start_dim=1) + + stats = torch.cat((pooling_mean, pooling_std), 1) + return stats + + +class ASTP(nn.Module): + """Attentive statistics pooling: Channel- and context-dependent + statistics pooling, first used in ECAPA_TDNN. + """ + + def __init__(self, in_dim, bottleneck_dim=128, global_context_att=False): + super(ASTP, self).__init__() + self.global_context_att = global_context_att + + # Use Conv1d with stride == 1 rather than Linear, then we don't + # need to transpose inputs. + if global_context_att: + self.linear1 = nn.Conv1d( + in_dim * 3, bottleneck_dim, kernel_size=1 + ) # equals W and b in the paper + else: + self.linear1 = nn.Conv1d( + in_dim, bottleneck_dim, kernel_size=1 + ) # equals W and b in the paper + self.linear2 = nn.Conv1d( + bottleneck_dim, in_dim, kernel_size=1 + ) # equals V and k in the paper + + def forward(self, x): + """ + x: a 3-dimensional tensor in tdnn-based architecture (B,F,T) + or a 4-dimensional tensor in resnet architecture (B,C,F,T) + 0-dim: batch-dimension, last-dim: time-dimension (frame-dimension) + """ + if len(x.shape) == 4: + x = x.reshape(x.shape[0], x.shape[1] * x.shape[2], x.shape[3]) + assert len(x.shape) == 3 + + if self.global_context_att: + context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x) + context_std = torch.sqrt(torch.var(x, dim=-1, keepdim=True) + 1e-10).expand_as(x) + x_in = torch.cat((x, context_mean, context_std), dim=1) + else: + x_in = x + + # DON'T use ReLU here! ReLU may be hard to converge. + alpha = torch.tanh(self.linear1(x_in)) # alpha = F.relu(self.linear1(x_in)) + alpha = torch.softmax(self.linear2(alpha), dim=2) + mean = torch.sum(alpha * x, dim=2) + var = torch.sum(alpha * (x**2), dim=2) - mean**2 + std = torch.sqrt(var.clamp(min=1e-10)) + return torch.cat([mean, std], dim=1) diff --git a/funasr/models/whisper_lid/eres2net/simple_avg.py b/funasr/models/whisper_lid/eres2net/simple_avg.py new file mode 100644 index 0000000000000000000000000000000000000000..9dbe7b3987b315eb76e09348a46e192302bb2494 --- /dev/null +++ b/funasr/models/whisper_lid/eres2net/simple_avg.py @@ -0,0 +1,18 @@ +import torch + +from funasr.models.encoder.abs_encoder import AbsEncoder +from funasr.modules.nets_utils import make_pad_mask + + +class SimpleAvg(AbsEncoder): + def __init__(self, feat_dim): + super(SimpleAvg, self).__init__() + self.feat_dim = feat_dim + + def forward(self, x, ilens): + mask = ~make_pad_mask(ilens, maxlen=x.shape[1]).to(x.device) + avg_x = (x * mask[:, :, None]).sum(1) / mask.sum(-1)[:, None] + return avg_x + + def output_size(self) -> int: + return self.feat_dim diff --git a/funasr/models/whisper_lid/lid_predictor.py b/funasr/models/whisper_lid/lid_predictor.py new file mode 100644 index 0000000000000000000000000000000000000000..d5b0d37ac4a5cd7235c7608d93e1b3c851a63e2d --- /dev/null +++ b/funasr/models/whisper_lid/lid_predictor.py @@ -0,0 +1,31 @@ +from funasr.register import tables +from funasr.models.whisper_lid.eres2net.ResNet import ( + ERes2Net, + BasicBlockERes2Net, + BasicBlockERes2Net_diff_AFF, +) + + +@tables.register("lid_predictor_classes", "LidPredictor") +class LidPredictor(ERes2Net): + def __init__( + self, + block=BasicBlockERes2Net, + block_fuse=BasicBlockERes2Net_diff_AFF, + num_blocks=[3, 4, 6, 3], + m_channels=32, + feat_dim=80, + embedding_size=192, + pooling_func="TSTP", + two_emb_layer=False, + ): + super(LidPredictor, self).__init__( + block=block, + block_fuse=block_fuse, + num_blocks=num_blocks, + m_channels=m_channels, + feat_dim=feat_dim, + embedding_size=embedding_size, + pooling_func=pooling_func, + two_emb_layer=two_emb_layer, + ) diff --git a/funasr/models/whisper_lid/model.py b/funasr/models/whisper_lid/model.py new file mode 100644 index 0000000000000000000000000000000000000000..02cd3730a886b7be6ee8dabb5de24d012e4599ba --- /dev/null +++ b/funasr/models/whisper_lid/model.py @@ -0,0 +1,702 @@ +import logging +from typing import Union, Dict, List, Tuple, Optional + +import time +import torch +import numpy as np +import torch.nn as nn +from torch.cuda.amp import autocast + +from funasr.losses.label_smoothing_loss import LabelSmoothingLoss +from funasr.models.ctc.ctc import CTC +from funasr.models.transformer.utils.add_sos_eos import add_sos_eos +from funasr.metrics.compute_acc import th_accuracy +from funasr.train_utils.device_funcs import force_gatherable +from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank +from funasr.utils import postprocess_utils +from funasr.utils.datadir_writer import DatadirWriter +from funasr.register import tables + + +@tables.register("model_classes", "OpenAIWhisperModel") +class OpenAIWhisperModel(nn.Module): + """CTC-attention hybrid Encoder-Decoder model""" + + def __init__( + self, + specaug: str = None, + specaug_conf: dict = None, + normalize: str = None, + normalize_conf: dict = None, + encoder: str = None, + encoder_conf: dict = None, + decoder: str = None, + decoder_conf: dict = None, + ctc: str = None, + ctc_conf: dict = None, + ctc_weight: float = 0.5, + interctc_weight: float = 0.0, + input_size: int = 80, + vocab_size: int = -1, + ignore_id: int = -1, + blank_id: int = 0, + sos: int = 1, + eos: int = 2, + lsm_weight: float = 0.0, + length_normalized_loss: bool = False, + report_cer: bool = True, + report_wer: bool = True, + sym_space: str = "", + sym_blank: str = "", + # extract_feats_in_collect_stats: bool = True, + share_embedding: bool = False, + # preencoder: Optional[AbsPreEncoder] = None, + # postencoder: Optional[AbsPostEncoder] = None, + **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() + if decoder is not None: + decoder_class = tables.decoder_classes.get(decoder) + decoder = decoder_class(decoder_conf) + if ctc_weight > 0.0: + + if ctc_conf is None: + ctc_conf = {} + + ctc = CTC(odim=vocab_size, encoder_output_size=encoder_output_size, **ctc_conf) + + self.blank_id = blank_id + self.sos = sos if sos is not None else vocab_size - 1 + self.eos = eos if eos is not None else vocab_size - 1 + self.vocab_size = vocab_size + self.ignore_id = ignore_id + self.ctc_weight = ctc_weight + self.specaug = specaug + self.normalize = normalize + self.encoder = encoder + + if not hasattr(self.encoder, "interctc_use_conditioning"): + self.encoder.interctc_use_conditioning = False + if self.encoder.interctc_use_conditioning: + self.encoder.conditioning_layer = torch.nn.Linear( + vocab_size, self.encoder.output_size() + ) + self.interctc_weight = interctc_weight + + # self.error_calculator = None + if ctc_weight == 1.0: + self.decoder = None + else: + self.decoder = decoder + + self.criterion_att = LabelSmoothingLoss( + size=vocab_size, + padding_idx=ignore_id, + smoothing=lsm_weight, + normalize_length=length_normalized_loss, + ) + # + # if report_cer or report_wer: + # self.error_calculator = ErrorCalculator( + # token_list, sym_space, sym_blank, report_cer, report_wer + # ) + # + self.error_calculator = None + if ctc_weight == 0.0: + self.ctc = None + else: + self.ctc = ctc + + self.share_embedding = share_embedding + if self.share_embedding: + self.decoder.embed = None + + self.length_normalized_loss = length_normalized_loss + 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]: + """Encoder + Decoder + Calc loss + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + text: (Batch, Length) + text_lengths: (Batch,) + """ + if len(text_lengths.size()) > 1: + text_lengths = text_lengths[:, 0] + if len(speech_lengths.size()) > 1: + speech_lengths = speech_lengths[:, 0] + + batch_size = speech.shape[0] + + # 1. Encoder + encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) + 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 = None, None, None, None + loss_ctc, cer_ctc = None, None + stats = dict() + + # decoder: CTC branch + if self.ctc_weight != 0.0: + loss_ctc, cer_ctc = self._calc_ctc_loss( + encoder_out, encoder_out_lens, text, text_lengths + ) + + # Collect CTC branch stats + stats["loss_ctc"] = loss_ctc.detach() if loss_ctc is not None else None + stats["cer_ctc"] = cer_ctc + + # Intermediate CTC (optional) + loss_interctc = 0.0 + if self.interctc_weight != 0.0 and intermediate_outs is not None: + for layer_idx, intermediate_out in intermediate_outs: + # we assume intermediate_out has the same length & padding + # as those of encoder_out + loss_ic, cer_ic = self._calc_ctc_loss( + intermediate_out, encoder_out_lens, text, text_lengths + ) + loss_interctc = loss_interctc + loss_ic + + # Collect Intermedaite CTC stats + stats["loss_interctc_layer{}".format(layer_idx)] = ( + loss_ic.detach() if loss_ic is not None else None + ) + stats["cer_interctc_layer{}".format(layer_idx)] = cer_ic + + loss_interctc = loss_interctc / len(intermediate_outs) + + # calculate whole encoder loss + loss_ctc = (1 - self.interctc_weight) * loss_ctc + self.interctc_weight * loss_interctc + + # decoder: Attention decoder branch + loss_att, acc_att, cer_att, wer_att = self._calc_att_loss( + encoder_out, encoder_out_lens, text, text_lengths + ) + + # 3. CTC-Att loss definition + if self.ctc_weight == 0.0: + loss = loss_att + elif self.ctc_weight == 1.0: + loss = loss_ctc + else: + loss = self.ctc_weight * loss_ctc + (1 - self.ctc_weight) * loss_att + + # Collect Attn branch stats + stats["loss_att"] = loss_att.detach() if loss_att is not None else None + stats["acc"] = acc_att + stats["cer"] = cer_att + stats["wer"] = wer_att + + # Collect total loss stats + 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 encode( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Frontend + Encoder. Note that this method is used by asr_inference.py + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + ind: int + """ + with autocast(False): + + # Data augmentation + if self.specaug is not None and self.training: + speech, speech_lengths = self.specaug(speech, speech_lengths) + + # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN + if self.normalize is not None: + speech, speech_lengths = self.normalize(speech, speech_lengths) + + # Forward encoder + # feats: (Batch, Length, Dim) + # -> encoder_out: (Batch, Length2, Dim2) + if self.encoder.interctc_use_conditioning: + encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths, ctc=self.ctc) + else: + encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths) + intermediate_outs = None + if isinstance(encoder_out, tuple): + intermediate_outs = encoder_out[1] + encoder_out = encoder_out[0] + + if intermediate_outs is not None: + return (encoder_out, intermediate_outs), encoder_out_lens + + return encoder_out, encoder_out_lens + + def _calc_att_loss( + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, + ): + 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_ctc_loss( + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, + ): + # Calc CTC loss + loss_ctc = self.ctc(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens) + + # Calc CER using CTC + cer_ctc = None + if not self.training and self.error_calculator is not None: + ys_hat = self.ctc.argmax(encoder_out).data + cer_ctc = self.error_calculator(ys_hat.cpu(), ys_pad.cpu(), is_ctc=True) + return loss_ctc, cer_ctc + + def init_beam_search( + self, + **kwargs, + ): + from funasr.models.transformer.search import BeamSearch + from funasr.models.transformer.scorers.ctc import CTCPrefixScorer + from funasr.models.transformer.scorers.length_bonus import LengthBonus + + # 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=self.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.5), + ctc=kwargs.get("decoding_ctc_weight", 0.5), + lm=kwargs.get("lm_weight", 0.0), + ngram=kwargs.get("ngram_weight", 0.0), + length_bonus=kwargs.get("penalty", 0.0), + ) + beam_search = BeamSearch( + beam_size=kwargs.get("beam_size", 10), + 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, + ): + + if kwargs.get("batch_size", 1) > 1: + raise NotImplementedError("batch decoding is not implemented") + + # init beamsearch + if self.beam_search is None: + logging.info("enable beam_search") + self.init_beam_search(**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"]) + # Encoder + encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) + if isinstance(encoder_out, tuple): + encoder_out = encoder_out[0] + + # c. Passed the encoder result and the beam search + nbest_hyps = self.beam_search( + x=encoder_out[0], + maxlenratio=kwargs.get("maxlenratio", 0.0), + minlenratio=kwargs.get("minlenratio", 0.0), + ) + + nbest_hyps = nbest_hyps[: self.nbest] + + results = [] + b, n, d = encoder_out.size() + for i in range(b): + + for nbest_idx, hyp in enumerate(nbest_hyps): + ibest_writer = None + if kwargs.get("output_dir") is not None: + if not hasattr(self, "writer"): + self.writer = DatadirWriter(kwargs.get("output_dir")) + ibest_writer = self.writer[f"{nbest_idx + 1}best_recog"] + + # remove sos/eos and get results + last_pos = -1 + if isinstance(hyp.yseq, list): + token_int = hyp.yseq[1:last_pos] + else: + token_int = hyp.yseq[1:last_pos].tolist() + + # remove blank symbol id, which is assumed to be 0 + token_int = list( + filter( + lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int + ) + ) + + # Change integer-ids to tokens + token = tokenizer.ids2tokens(token_int) + text = tokenizer.tokens2text(token) + + text_postprocessed, _ = postprocess_utils.sentence_postprocess(token) + result_i = {"key": key[i], "token": token, "text": text_postprocessed} + results.append(result_i) + + if ibest_writer is not None: + ibest_writer["token"][key[i]] = " ".join(token) + ibest_writer["text"][key[i]] = text_postprocessed + + return results, meta_data + + +@tables.register("model_classes", "OpenAIWhisperLIDModel") +class OpenAIWhisperLIDModel(nn.Module): + """WhisperEncoder and EResNet based LID Model""" + + def __init__( + self, + vocab_size: int, + specaug: str = None, + specaug_conf: dict = None, + encoder: str = None, + encoder_conf: dict = None, + lid_predictor: str = None, + lid_predictor_conf: dict = None, + proj_dim: int = None, + clip_frames: int = None, + random_clip: bool = False, + **kwargs, + ): + super().__init__() + if specaug is not None: + specaug_class = tables.specaug_classes.get(specaug) + specaug = specaug_class(**specaug_conf) + encoder_class = tables.encoder_classes.get(encoder) + encoder = encoder_class(**encoder_conf) + lid_predictor_class = tables.lid_predictor_classes.get(lid_predictor) + lid_predictor = lid_predictor_class(**lid_predictor_conf) + if encoder.output_size() != proj_dim: + self.proj_layer = torch.nn.Linear(encoder.output_size(), proj_dim) + else: + self.proj_layer = None + self.output_layer = torch.nn.Linear(lid_predictor.output_size(), vocab_size) + self.criterion_lid = LabelSmoothingLoss( + size=vocab_size, + padding_idx=-1, + smoothing=0.0, + normalize_length=False, + ) + + self.specaug = specaug + self.encoder = encoder + self.lid_predictor = lid_predictor + self.clip_frames = clip_frames + self.random_clip = random_clip + self.normalize = None + self.beam_search = None + if not hasattr(self.encoder, "interctc_use_conditioning"): + self.encoder.interctc_use_conditioning = False + + def forward( + self, + speech: torch.Tensor, # may be padding + speech_lengths: torch.Tensor, # actual length + lid: torch.Tensor, # lid label, (batch_size, 1) + lid_lengths: torch.Tensor, + ): + assert lid.shape[1] == 1 + batch_size = speech.shape[0] + encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) + + # re-generate encoder_out + if self.clip_frames is None: + reduced_encoder_out = ( + torch.zeros(batch_size, encoder_out_lens.max(), encoder_out.shape[-1]) + .to(encoder_out.dtype) + .to(encoder_out.device) + ) + for i, enc_length in enumerate(encoder_out_lens): + reduced_encoder_out[i, :enc_length] = encoder_out[i, :enc_length] + else: + reduced_encoder_out = ( + torch.zeros(batch_size, self.clip_frames, encoder_out.shape[-1]) + .to(encoder_out.dtype) + .to(encoder_out.device) + ) + if self.random_clip: + for i, enc_length in enumerate(encoder_out_lens): + if enc_length <= self.clip_frames: + reduced_encoder_out[i, :enc_length] = encoder_out[i, :enc_length] + encoder_out_lens[i] = enc_length + else: + max_start_index = enc_length.item() - self.clip_frames + start_index = np.random.randint(0, max_start_index + 1) + reduced_encoder_out[i, : self.clip_frames] = encoder_out[ + i, start_index : start_index + self.clip_frames + ] + encoder_out_lens[i] = self.clip_frames + else: + for i, enc_length in enumerate(encoder_out_lens): + enc_length = self.clip_frames if enc_length >= self.clip_frames else enc_length + reduced_encoder_out[i, :enc_length] = encoder_out[i, :enc_length] + encoder_out_lens[i] = enc_length + if self.proj_layer is not None: + reduced_encoder_out = self.proj_layer(reduced_encoder_out) + lid_output = self.lid_predictor(reduced_encoder_out, encoder_out_lens) # (B, D) + lid_logits = self.output_layer(lid_output) # (B, num_classes) + loss = self.criterion_lid(lid_logits[:, None, :], lid) + with torch.no_grad(): + _, predicted_lid = torch.max(lid_logits, 1) + correct = (predicted_lid == lid[:, 0]).sum().item() + lid_acc = correct * 1.0 / lid_logits.shape[0] + stats = dict() + stats["batch_size"] = batch_size + stats["loss"] = torch.clone(loss.detach()) + stats["acc"] = lid_acc + stats["token_length"] = speech_lengths.max() + loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) + return loss, stats, weight + + def encode( + self, speech: torch.Tensor, speech_lengths: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Frontend + Encoder. Note that this method is used by asr_inference.py + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + """ + with autocast(False): + + # Data augmentation + if self.specaug is not None and self.training: + speech = speech.permute(0, 2, 1) + # suit for whisper padding + padded_speech_lengths = torch.ones_like(speech_lengths) * speech.shape[1] + speech, padded_speech_lengths = self.specaug(speech, padded_speech_lengths) + speech = speech.permute(0, 2, 1) + + # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN + if self.normalize is not None: + speech, speech_lengths = self.normalize(speech, speech_lengths) + + # Forward encoder + # feats: (Batch, Length, Dim) + # -> encoder_out: (Batch, Length2, Dim2) + if self.encoder.interctc_use_conditioning: + encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths, ctc=self.ctc) + else: + encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths) + intermediate_outs = None + if isinstance(encoder_out, tuple): + intermediate_outs = encoder_out[1] + encoder_out = encoder_out[0] + + if intermediate_outs is not None: + return (encoder_out, intermediate_outs), encoder_out_lens + + return encoder_out, encoder_out_lens + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): + + if kwargs.get("batch_size", 1) > 1: + raise NotImplementedError("batch decoding is not implemented") + + 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"]) + # Encoder + enc, enc_out_lens = self.encode(speech, speech_lengths) + + inference_clip_length = kwargs.get("inference_clip_length", None) + if self.clip_frames is not None: + if inference_clip_length is None: + reduced_enc = ( + torch.zeros(enc.shape[0], self.clip_frames, enc.shape[-1]) + .to(enc.dtype) + .to(enc.device) + ) + for i, enc_length in enumerate(enc_out_lens): + enc_length = self.clip_frames if enc_length >= self.clip_frames else enc_length + reduced_enc[i, :enc_length] = enc[i, :enc_length] + enc_out_lens[i] = enc_length + else: + assert inference_clip_length > 0, "inference_clip_length must be larger than 0" + reduced_enc = ( + torch.zeros(enc.shape[0], inference_clip_length, enc.shape[-1]) + .to(enc.dtype) + .to(enc.device) + ) + for i, enc_length in enumerate(enc_out_lens): + enc_length = ( + inference_clip_length if enc_length >= inference_clip_length else enc_length + ) + reduced_enc[i, :enc_length] = enc[i, :enc_length] + enc_out_lens[i] = enc_length + else: + reduced_enc = ( + torch.zeros(enc.shape[0], enc_out_lens.max(), enc.shape[-1]) + .to(enc.dtype) + .to(enc.device) + ) + for i, enc_length in enumerate(enc_out_lens): + reduced_enc[i, :enc_length] = enc[i, :enc_length] + + if self.proj_layer is not None: + reduced_enc = self.proj_layer(reduced_enc) + lid_output = self.lid_predictor(reduced_enc, enc_out_lens) # (B, D) + lid_logits = self.output_layer(lid_output) # (B, num_classes) + + _, predicted_lid_index = torch.max(lid_logits, 1) + predicted_lid = tokenizer.ids2tokens([predicted_lid_index[0].cpu()])[0] + + if kwargs.get("output_dir") is not None: + if not hasattr(self, "writer"): + self.writer = DatadirWriter(kwargs.get("output_dir")) + lid_writer = self.writer["lid"] + lid_writer[key[0]] = predicted_lid + + results = [{"key": key[0], "lid": predicted_lid}] + + return results, meta_data diff --git a/funasr/models/xvector/__init__.py b/funasr/models/xvector/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/models/xvector/e2e_sv.py b/funasr/models/xvector/e2e_sv.py new file mode 100644 index 0000000000000000000000000000000000000000..56c9e3b10a2350b3e2af2a1481af5ba5a2e74d6b --- /dev/null +++ b/funasr/models/xvector/e2e_sv.py @@ -0,0 +1,263 @@ +""" +Author: Speech Lab, Alibaba Group, China +""" + +import logging +from contextlib import contextmanager +from distutils.version import LooseVersion +from typing import Dict +from typing import List +from typing import Optional +from typing import Tuple +from typing import Union + +import torch + +from funasr.layers.abs_normalize import AbsNormalize +from funasr.losses.label_smoothing_loss import ( + LabelSmoothingLoss, # noqa: H301 +) +from funasr.models.ctc import CTC +from funasr.models.decoder.abs_decoder import AbsDecoder +from funasr.models.encoder.abs_encoder import AbsEncoder +from funasr.frontends.abs_frontend import AbsFrontend +from funasr.models.postencoder.abs_postencoder import AbsPostEncoder +from funasr.models.preencoder.abs_preencoder import AbsPreEncoder +from funasr.models.specaug.abs_specaug import AbsSpecAug +from funasr.models.transformer.utils.add_sos_eos import add_sos_eos +from funasr.metrics import ErrorCalculator +from funasr.metrics.compute_acc import th_accuracy +from funasr.train_utils.device_funcs import force_gatherable +from funasr.models.base_model import FunASRModel + +if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): + from torch.cuda.amp import autocast +else: + # Nothing to do if torch<1.6.0 + @contextmanager + def autocast(enabled=True): + yield + + +class ESPnetSVModel(FunASRModel): + """CTC-attention hybrid Encoder-Decoder model""" + + def __init__( + self, + vocab_size: int, + token_list: Union[Tuple[str, ...], List[str]], + frontend: Optional[AbsFrontend], + specaug: Optional[AbsSpecAug], + normalize: Optional[AbsNormalize], + preencoder: Optional[AbsPreEncoder], + encoder: AbsEncoder, + postencoder: Optional[AbsPostEncoder], + pooling_layer: torch.nn.Module, + decoder: AbsDecoder, + ): + + super().__init__() + # note that eos is the same as sos (equivalent ID) + self.vocab_size = vocab_size + self.token_list = token_list.copy() + + self.frontend = frontend + self.specaug = specaug + self.normalize = normalize + self.preencoder = preencoder + self.postencoder = postencoder + self.encoder = encoder + self.pooling_layer = pooling_layer + self.decoder = decoder + + def forward( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + text: torch.Tensor, + text_lengths: torch.Tensor, + ) -> 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,) + """ + assert text_lengths.dim() == 1, text_lengths.shape + # Check that batch_size is unified + assert ( + speech.shape[0] == speech_lengths.shape[0] == text.shape[0] == text_lengths.shape[0] + ), (speech.shape, speech_lengths.shape, text.shape, text_lengths.shape) + batch_size = speech.shape[0] + + # for data-parallel + text = text[:, : text_lengths.max()] + + # 1. Encoder + encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) + 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 = None, None, None, None + loss_ctc, cer_ctc = None, None + loss_transducer, cer_transducer, wer_transducer = None, None, None + stats = dict() + + # 1. CTC branch + if self.ctc_weight != 0.0: + loss_ctc, cer_ctc = self._calc_ctc_loss( + encoder_out, encoder_out_lens, text, text_lengths + ) + + # Collect CTC branch stats + stats["loss_ctc"] = loss_ctc.detach() if loss_ctc is not None else None + stats["cer_ctc"] = cer_ctc + + # Intermediate CTC (optional) + loss_interctc = 0.0 + if self.interctc_weight != 0.0 and intermediate_outs is not None: + for layer_idx, intermediate_out in intermediate_outs: + # we assume intermediate_out has the same length & padding + # as those of encoder_out + loss_ic, cer_ic = self._calc_ctc_loss( + intermediate_out, encoder_out_lens, text, text_lengths + ) + loss_interctc = loss_interctc + loss_ic + + # Collect Intermedaite CTC stats + stats["loss_interctc_layer{}".format(layer_idx)] = ( + loss_ic.detach() if loss_ic is not None else None + ) + stats["cer_interctc_layer{}".format(layer_idx)] = cer_ic + + loss_interctc = loss_interctc / len(intermediate_outs) + + # calculate whole encoder loss + loss_ctc = (1 - self.interctc_weight) * loss_ctc + self.interctc_weight * loss_interctc + + if self.use_transducer_decoder: + # 2a. Transducer decoder branch + ( + loss_transducer, + cer_transducer, + wer_transducer, + ) = self._calc_transducer_loss( + encoder_out, + encoder_out_lens, + text, + ) + + if loss_ctc is not None: + loss = loss_transducer + (self.ctc_weight * loss_ctc) + else: + loss = loss_transducer + + # Collect Transducer branch stats + stats["loss_transducer"] = ( + loss_transducer.detach() if loss_transducer is not None else None + ) + stats["cer_transducer"] = cer_transducer + stats["wer_transducer"] = wer_transducer + + else: + # 2b. Attention decoder branch + if self.ctc_weight != 1.0: + loss_att, acc_att, cer_att, wer_att = self._calc_att_loss( + encoder_out, encoder_out_lens, text, text_lengths + ) + + # 3. CTC-Att loss definition + if self.ctc_weight == 0.0: + loss = loss_att + elif self.ctc_weight == 1.0: + loss = loss_ctc + else: + loss = self.ctc_weight * loss_ctc + (1 - self.ctc_weight) * loss_att + + # Collect Attn branch stats + stats["loss_att"] = loss_att.detach() if loss_att is not None else None + stats["acc"] = acc_att + stats["cer"] = cer_att + stats["wer"] = wer_att + + # Collect total loss stats + stats["loss"] = torch.clone(loss.detach()) + + # force_gatherable: to-device and to-tensor if scalar for DataParallel + 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 + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Frontend + Encoder. Note that this method is used by asr_inference.py + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + """ + with autocast(False): + # 1. Extract feats + feats, feats_lengths = self._extract_feats(speech, speech_lengths) + + # 2. Data augmentation + if self.specaug is not None and self.training: + feats, feats_lengths = self.specaug(feats, feats_lengths) + + # 3. Normalization for feature: e.g. Global-CMVN, Utterance-CMVN + if self.normalize is not None: + feats, feats_lengths = self.normalize(feats, feats_lengths) + + # Pre-encoder, e.g. used for raw input data + if self.preencoder is not None: + feats, feats_lengths = self.preencoder(feats, feats_lengths) + + # 4. Forward encoder + # feats: (Batch, Length, Dim) -> (Batch, Channel, Length2, Dim2) + encoder_out, encoder_out_lens = self.encoder(feats, feats_lengths) + + # Post-encoder, e.g. NLU + if self.postencoder is not None: + encoder_out, encoder_out_lens = self.postencoder(encoder_out, encoder_out_lens) + + return encoder_out, encoder_out_lens + + def _extract_feats( + self, speech: torch.Tensor, speech_lengths: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + assert speech_lengths.dim() == 1, speech_lengths.shape + + # for data-parallel + speech = speech[:, : speech_lengths.max()] + + if self.frontend is not None: + # Frontend + # e.g. STFT and Feature extract + # data_loader may send time-domain signal in this case + # speech (Batch, NSamples) -> feats: (Batch, NFrames, Dim) + feats, feats_lengths = self.frontend(speech, speech_lengths) + else: + # No frontend and no feature extract + feats, feats_lengths = speech, speech_lengths + return feats, feats_lengths diff --git a/funasr/optimizers/__init__.py b/funasr/optimizers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..66e2e9d040d286fbd6468ac922ade1c3da286d6f --- /dev/null +++ b/funasr/optimizers/__init__.py @@ -0,0 +1,17 @@ +import torch +from funasr.optimizers.fairseq_adam import FairseqAdam +from funasr.optimizers.sgd import SGD + +optim_classes = dict( + adam=torch.optim.Adam, + fairseq_adam=FairseqAdam, + adamw=torch.optim.AdamW, + sgd=SGD, + adadelta=torch.optim.Adadelta, + adagrad=torch.optim.Adagrad, + adamax=torch.optim.Adamax, + asgd=torch.optim.ASGD, + lbfgs=torch.optim.LBFGS, + rmsprop=torch.optim.RMSprop, + rprop=torch.optim.Rprop, +) diff --git a/funasr/optimizers/fairseq_adam.py b/funasr/optimizers/fairseq_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..32fae6a2582c908f157838e4bb0e6afddc200cdf --- /dev/null +++ b/funasr/optimizers/fairseq_adam.py @@ -0,0 +1,144 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + + +import math + +import torch +import torch.optim + + +class FairseqAdam(torch.optim.Optimizer): + r"""Implements Adam algorithm. + + This implementation is modified from torch.optim.Adam based on: + `Fixed Weight Decay Regularization in Adam` + (see https://arxiv.org/abs/1711.05101) + + It has been proposed in `Adam: A Method for Stochastic Optimization`_. + + Args: + params (iterable): iterable of parameters to optimize or dicts defining + parameter groups + lr (float, optional): learning rate (default: 1e-3) + betas (Tuple[float, float], optional): coefficients used for computing + running averages of gradient and its square (default: (0.9, 0.999)) + eps (float, optional): term added to the denominator to improve + numerical stability (default: 1e-8) + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + amsgrad (boolean, optional): whether to use the AMSGrad variant of this + algorithm from the paper `On the Convergence of Adam and Beyond`_ + + .. _Adam\: A Method for Stochastic Optimization: + https://arxiv.org/abs/1412.6980 + .. _On the Convergence of Adam and Beyond: + https://openreview.net/forum?id=ryQu7f-RZ + """ + + def __init__( + self, + params, + lr=1e-3, + adam_betas=(0.9, 0.999), + adam_eps=1e-8, + weight_decay=0, + amsgrad=False, + ): + defaults = dict( + lr=lr, betas=adam_betas, eps=adam_eps, weight_decay=weight_decay, amsgrad=amsgrad + ) + super(FairseqAdam, self).__init__(params, defaults) + self.optimizer_lr = lr + + @property + def supports_memory_efficient_fp16(self): + return True + + @property + def supports_flat_params(self): + return True + + def step(self, closure=None): + """Performs a single optimization step. + + Args: + closure (callable, optional): A closure that reevaluates the model + and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad.data + if grad.dtype in {torch.float16, torch.bfloat16}: + grad = grad.float() + if grad.is_sparse: + raise RuntimeError( + "Adam does not support sparse gradients, please consider SparseAdam instead" + ) + amsgrad = group.get("amsgrad", False) + + p_data_fp32 = p.data + if p.data.dtype in {torch.float16, torch.bfloat16}: + p_data_fp32 = p_data_fp32.float() + + state = self.state[p] + + # State initialization + if len(state) == 0: + state["step"] = 0 + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like(p_data_fp32) + # Exponential moving average of squared gradient values + state["exp_avg_sq"] = torch.zeros_like(p_data_fp32) + if amsgrad: + # Maintains max of all exp. moving avg. of sq. grad. values + state["max_exp_avg_sq"] = torch.zeros_like(p_data_fp32) + else: + state["exp_avg"] = state["exp_avg"].to(p_data_fp32) + state["exp_avg_sq"] = state["exp_avg_sq"].to(p_data_fp32) + if amsgrad: + state["max_exp_avg_sq"] = state["max_exp_avg_sq"].to(p_data_fp32) + + exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"] + if amsgrad: + max_exp_avg_sq = state["max_exp_avg_sq"] + beta1, beta2 = group["betas"] + + state["step"] += 1 + + # Decay the first and second moment running average coefficient + exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) + if amsgrad: + # Maintains the maximum of all 2nd moment running avg. till now + torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq) + # Use the max. for normalizing running avg. of gradient + denom = max_exp_avg_sq.sqrt().add_(group["eps"]) + else: + denom = exp_avg_sq.sqrt().add_(group["eps"]) + + bias_correction1 = 1 - beta1 ** state["step"] + bias_correction2 = 1 - beta2 ** state["step"] + step_size = group["lr"] * math.sqrt(bias_correction2) / bias_correction1 + + if group["weight_decay"] != 0: + p_data_fp32.add_(p_data_fp32, alpha=-group["weight_decay"] * group["lr"]) + + p_data_fp32.addcdiv_(exp_avg, denom, value=-step_size) + + if p.data.dtype in {torch.float16, torch.bfloat16}: + p.data.copy_(p_data_fp32) + + return loss + + def set_lr(self, lr): + """Set the learning rate.""" + for param_group in self.param_groups: + param_group["lr"] = lr diff --git a/funasr/optimizers/sgd.py b/funasr/optimizers/sgd.py new file mode 100644 index 0000000000000000000000000000000000000000..fb7a3df01ed428038361626ab1a8e40405bedb14 --- /dev/null +++ b/funasr/optimizers/sgd.py @@ -0,0 +1,30 @@ +import torch + + +class SGD(torch.optim.SGD): + """Thin inheritance of torch.optim.SGD to bind the required arguments, 'lr' + + Note that + the arguments of the optimizer invoked by AbsTask.main() + must have default value except for 'param'. + + I can't understand why only SGD.lr doesn't have the default value. + """ + + def __init__( + self, + params, + lr: float = 0.1, + momentum: float = 0.0, + dampening: float = 0.0, + weight_decay: float = 0.0, + nesterov: bool = False, + ): + super().__init__( + params, + lr=lr, + momentum=momentum, + dampening=dampening, + weight_decay=weight_decay, + nesterov=nesterov, + ) diff --git a/funasr/register.py b/funasr/register.py new file mode 100644 index 0000000000000000000000000000000000000000..a8b1fd333a3239b67365d8c7763cd3d4c71735a1 --- /dev/null +++ b/funasr/register.py @@ -0,0 +1,87 @@ +import logging +import inspect +from dataclasses import dataclass +import re + + +@dataclass +class RegisterTables: + """Registry system for classes.""" + + model_classes = {} + frontend_classes = {} + specaug_classes = {} + normalize_classes = {} + encoder_classes = {} + decoder_classes = {} + joint_network_classes = {} + predictor_classes = {} + stride_conv_classes = {} + tokenizer_classes = {} + dataloader_classes = {} + batch_sampler_classes = {} + dataset_classes = {} + index_ds_classes = {} + + def print(self, key: str = None) -> None: + """Print registered classes.""" + print("\ntables: \n") + fields = vars(self) + headers = ["register name", "class name", "class location"] + for classes_key, classes_dict in fields.items(): + if classes_key.endswith("_meta") and (key is None or key in classes_key): + print(f"----------- ** {classes_key.replace('_meta', '')} ** --------------") + metas = [] + for register_key, meta in classes_dict.items(): + metas.append(meta) + metas.sort(key=lambda x: x[0]) + data = [headers] + metas + col_widths = [max(len(str(item)) for item in col) for col in zip(*data)] + + for row in data: + print( + "| " + + " | ".join(str(item).ljust(width) for item, width in zip(row, col_widths)) + + " |" + ) + print("\n") + + def register(self, register_tables_key: str, key: str = None) -> callable: + """Decorator to register a class.""" + + def decorator(target_class): + if not hasattr(self, register_tables_key): + setattr(self, register_tables_key, {}) + logging.debug(f"New registry table added: {register_tables_key}") + + registry = getattr(self, register_tables_key) + registry_key = key if key is not None else target_class.__name__ + + if registry_key in registry: + logging.debug( + f"Key {registry_key} already exists in {register_tables_key}, re-register" + ) + + registry[registry_key] = target_class + + register_tables_key_meta = register_tables_key + "_meta" + if not hasattr(self, register_tables_key_meta): + setattr(self, register_tables_key_meta, {}) + registry_meta = getattr(self, register_tables_key_meta) + + class_file = inspect.getfile(target_class) + class_line = inspect.getsourcelines(target_class)[1] + pattern = r"^.+/funasr/" + class_file = re.sub(pattern, "funasr/", class_file) + meta_data = [ + registry_key, + target_class.__name__, + f"{class_file}:{class_line}", + ] + registry_meta[registry_key] = meta_data + return target_class + + return decorator + + +tables = RegisterTables() diff --git a/funasr/schedulers/__init__.py b/funasr/schedulers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..39f8c0ece80687f010bc8f79e0d2f47f59662c0b --- /dev/null +++ b/funasr/schedulers/__init__.py @@ -0,0 +1,25 @@ +import torch +import torch.multiprocessing +import torch.nn +import torch.optim + +from funasr.schedulers.noam_lr import NoamLR +from funasr.schedulers.tri_stage_scheduler import TriStageLR +from funasr.schedulers.warmup_lr import WarmupLR +from funasr.schedulers.lambdalr_cus import CustomLambdaLR + +scheduler_classes = dict( + ReduceLROnPlateau=torch.optim.lr_scheduler.ReduceLROnPlateau, + lambdalr=torch.optim.lr_scheduler.LambdaLR, + steplr=torch.optim.lr_scheduler.StepLR, + multisteplr=torch.optim.lr_scheduler.MultiStepLR, + exponentiallr=torch.optim.lr_scheduler.ExponentialLR, + CosineAnnealingLR=torch.optim.lr_scheduler.CosineAnnealingLR, + noamlr=NoamLR, + warmuplr=WarmupLR, + tri_stage=TriStageLR, + cycliclr=torch.optim.lr_scheduler.CyclicLR, + onecyclelr=torch.optim.lr_scheduler.OneCycleLR, + CosineAnnealingWarmRestarts=torch.optim.lr_scheduler.CosineAnnealingWarmRestarts, + custom_lambdalr=CustomLambdaLR, +) diff --git a/funasr/schedulers/abs_scheduler.py b/funasr/schedulers/abs_scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..7395f259c3e563f560ea99080e00dac61967b421 --- /dev/null +++ b/funasr/schedulers/abs_scheduler.py @@ -0,0 +1,84 @@ +from abc import ABC +from abc import abstractmethod + +import torch.optim.lr_scheduler as L + + +class AbsScheduler(ABC): + @abstractmethod + def step(self, epoch: int = None): + pass + + @abstractmethod + def state_dict(self): + pass + + @abstractmethod + def load_state_dict(self, state): + pass + + +# If you need to define custom scheduler, please inherit these classes +class AbsBatchStepScheduler(AbsScheduler): + @abstractmethod + def step(self, epoch: int = None): + pass + + @abstractmethod + def state_dict(self): + pass + + @abstractmethod + def load_state_dict(self, state): + pass + + +class AbsEpochStepScheduler(AbsScheduler): + @abstractmethod + def step(self, epoch: int = None): + pass + + @abstractmethod + def state_dict(self): + pass + + @abstractmethod + def load_state_dict(self, state): + pass + + +class AbsValEpochStepScheduler(AbsEpochStepScheduler): + @abstractmethod + def step(self, val, epoch: int = None): + pass + + @abstractmethod + def state_dict(self): + pass + + @abstractmethod + def load_state_dict(self, state): + pass + + +# Create alias type to check the type +# Note(kamo): Currently PyTorch doesn't provide the base class +# to judge these classes. +AbsValEpochStepScheduler.register(L.ReduceLROnPlateau) +for s in [ + L.ReduceLROnPlateau, + L.LambdaLR, + L.StepLR, + L.MultiStepLR, + L.MultiStepLR, + L.ExponentialLR, + L.CosineAnnealingLR, +]: + AbsEpochStepScheduler.register(s) + +AbsBatchStepScheduler.register(L.CyclicLR) +for s in [ + L.OneCycleLR, + L.CosineAnnealingWarmRestarts, +]: + AbsBatchStepScheduler.register(s) diff --git a/funasr/schedulers/lambdalr_cus.py b/funasr/schedulers/lambdalr_cus.py new file mode 100644 index 0000000000000000000000000000000000000000..e3bb1fb47581a88850d1e29b5d26d20fe5ae7538 --- /dev/null +++ b/funasr/schedulers/lambdalr_cus.py @@ -0,0 +1,41 @@ +import torch +from torch.optim.lr_scheduler import _LRScheduler + + +# class CustomLambdaLR(_LRScheduler): +# def __init__(self, optimizer, warmup_steps, last_epoch=-1): +# self.warmup_steps = warmup_steps +# super().__init__(optimizer, last_epoch) +# +# def get_lr(self): +# if self.last_epoch < self.warmup_steps: +# return [ +# base_lr * min(self.last_epoch / self.warmup_steps, 1) for base_lr in self.base_lrs +# ] +# else: +# return [base_lr for base_lr in self.base_lrs] + + +class CustomLambdaLR(_LRScheduler): + def __init__( + self, + optimizer, + warmup_steps: int = 25000, + total_steps: int = 500000, + last_epoch=-1, + verbose=False, + ): + self.warmup_steps = warmup_steps + self.total_steps = total_steps + super().__init__(optimizer, last_epoch, verbose) + + def get_lr(self): + + step = self.last_epoch + 1 + if step < self.warmup_steps: + lr_scale = step / self.warmup_steps + else: + lr_scale = max( + 0.0, 1 - (step - self.warmup_steps) / (self.total_steps - self.warmup_steps) + ) + return [base_lr * lr_scale for base_lr in self.base_lrs] diff --git a/funasr/schedulers/noam_lr.py b/funasr/schedulers/noam_lr.py new file mode 100644 index 0000000000000000000000000000000000000000..b33b4bde610ac98f35d85d3092ad729f2512e8ee --- /dev/null +++ b/funasr/schedulers/noam_lr.py @@ -0,0 +1,62 @@ +"""Noam learning rate scheduler module.""" + +from typing import Union +import warnings + +import torch +from torch.optim.lr_scheduler import _LRScheduler + +from funasr.schedulers.abs_scheduler import AbsBatchStepScheduler + + +class NoamLR(_LRScheduler, AbsBatchStepScheduler): + """The LR scheduler proposed by Noam + + Ref: + "Attention Is All You Need", https://arxiv.org/pdf/1706.03762.pdf + + FIXME(kamo): PyTorch doesn't provide _LRScheduler as public class, + thus the behaviour isn't guaranteed at forward PyTorch version. + + NOTE(kamo): The "model_size" in original implementation is derived from + the model, but in this implementation, this parameter is a constant value. + You need to change it if the model is changed. + + """ + + def __init__( + self, + optimizer: torch.optim.Optimizer, + model_size: Union[int, float] = 320, + warmup_steps: Union[int, float] = 25000, + last_epoch: int = -1, + ): + self.model_size = model_size + self.warmup_steps = warmup_steps + + lr = list(optimizer.param_groups)[0]["lr"] + new_lr = self.lr_for_WarmupLR(lr) + warnings.warn( + f"NoamLR is deprecated. " + f"Use WarmupLR(warmup_steps={warmup_steps}) with Optimizer(lr={new_lr})", + ) + + # __init__() must be invoked before setting field + # because step() is also invoked in __init__() + super().__init__(optimizer, last_epoch) + + def lr_for_WarmupLR(self, lr: float) -> float: + return lr / self.model_size**0.5 / self.warmup_steps**0.5 + + def __repr__(self): + return ( + f"{self.__class__.__name__}(model_size={self.model_size}, " + f"warmup_steps={self.warmup_steps})" + ) + + def get_lr(self): + step_num = self.last_epoch + 1 + return [ + lr * self.model_size**-0.5 * min(step_num**-0.5, step_num * self.warmup_steps**-1.5) + for lr in self.base_lrs + ] diff --git a/funasr/schedulers/tri_stage_scheduler.py b/funasr/schedulers/tri_stage_scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..7fc13b2ffe0c72f17f4415578b468d787ba4edd2 --- /dev/null +++ b/funasr/schedulers/tri_stage_scheduler.py @@ -0,0 +1,104 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import math +from typing import Optional, List + +import torch +from torch.optim.lr_scheduler import _LRScheduler + +from funasr.schedulers.abs_scheduler import AbsBatchStepScheduler + + +class TriStageLR(_LRScheduler, AbsBatchStepScheduler): + def __init__( + self, + optimizer: torch.optim.Optimizer, + last_epoch: int = -1, + phase_ratio: Optional[List[float]] = None, + init_lr_scale: float = 0.01, + final_lr_scale: float = 0.01, + ): + self.optimizer = optimizer + self.last_epoch = last_epoch + self.phase_ratio = phase_ratio + self.init_lr_scale = init_lr_scale + self.final_lr_scale = final_lr_scale + self.optimizer_lr = self.optimizer.defaults["lr"] + + def init_tri_stage_scheudler(self, max_update): + self.max_update = max_update + self.peak_lr = self.optimizer_lr + self.init_lr = self.init_lr_scale * self.optimizer_lr + self.final_lr = self.final_lr_scale * self.optimizer_lr + + assert self.max_update > 0 + assert sum(self.phase_ratio) == 1, "phase ratios must add up to 1" + assert len(self.phase_ratio) == 3 + self.warmup_steps = int(self.max_update * self.phase_ratio[0]) + self.hold_steps = int(self.max_update * self.phase_ratio[1]) + self.decay_steps = int(self.max_update * self.phase_ratio[2]) + + self.warmup_rate = ( + (self.peak_lr - self.init_lr) / self.warmup_steps if self.warmup_steps != 0 else 0 + ) + self.decay_factor = -math.log(self.final_lr_scale) / self.decay_steps + + # initial learning rate + self.lr = self.init_lr + + # __init__() must be invoked before setting field + # because step() is also invoked in __init__() + self.set_optimizer_lr(self.lr) + super().__init__(self.optimizer, self.last_epoch) + + def _decide_stage(self, update_step): + """ + return stage, and the corresponding steps within the current stage + """ + if update_step < self.warmup_steps: + # warmup state + return 0, update_step + + offset = self.warmup_steps + + if update_step < offset + self.hold_steps: + # hold stage + return 1, update_step - offset + + offset += self.hold_steps + + if update_step <= offset + self.decay_steps: + # decay stage + return 2, update_step - offset + + offset += self.decay_steps + + # still here ? constant lr stage + return 3, update_step - offset + + def step_update(self, num_updates): + """Update the learning rate after each update.""" + stage, steps_in_stage = self._decide_stage(num_updates) + if stage == 0: + self.lr = self.init_lr + self.warmup_rate * steps_in_stage + elif stage == 1: + self.lr = self.peak_lr + elif stage == 2: + self.lr = self.peak_lr * math.exp(-self.decay_factor * steps_in_stage) + elif stage == 3: + self.lr = self.final_lr + else: + raise ValueError("Undefined stage") + self.set_optimizer_lr(self.lr) + + def set_optimizer_lr(self, lr): + for param_group in self.optimizer.param_groups: + param_group["lr"] = lr + + def get_lr(self): + step_num = self.last_epoch + 1 + self.step_update(step_num) + return [self.lr] diff --git a/funasr/schedulers/warmup_lr.py b/funasr/schedulers/warmup_lr.py new file mode 100644 index 0000000000000000000000000000000000000000..847b53ad705d1b786e8572e87bdf5125d8148335 --- /dev/null +++ b/funasr/schedulers/warmup_lr.py @@ -0,0 +1,47 @@ +"""Warm up learning rate scheduler module.""" + +from typing import Union + +import torch +from torch.optim.lr_scheduler import _LRScheduler + +from funasr.schedulers.abs_scheduler import AbsBatchStepScheduler + + +class WarmupLR(_LRScheduler, AbsBatchStepScheduler): + """The WarmupLR scheduler + + This scheduler is almost same as NoamLR Scheduler except for following difference: + + NoamLR: + lr = optimizer.lr * model_size ** -0.5 + * min(step ** -0.5, step * warmup_step ** -1.5) + WarmupLR: + lr = optimizer.lr * warmup_step ** 0.5 + * min(step ** -0.5, step * warmup_step ** -1.5) + + Note that the maximum lr equals to optimizer.lr in this scheduler. + + """ + + def __init__( + self, + optimizer: torch.optim.Optimizer, + warmup_steps: Union[int, float] = 25000, + last_epoch: int = -1, + ): + self.warmup_steps = warmup_steps + + # __init__() must be invoked before setting field + # because step() is also invoked in __init__() + super().__init__(optimizer, last_epoch) + + def __repr__(self): + return f"{self.__class__.__name__}(warmup_steps={self.warmup_steps})" + + def get_lr(self): + step_num = self.last_epoch + 1 + return [ + lr * self.warmup_steps**0.5 * min(step_num**-0.5, step_num * self.warmup_steps**-1.5) + for lr in self.base_lrs + ] diff --git a/funasr/tokenizer/__init__.py b/funasr/tokenizer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/tokenizer/abs_tokenizer.py b/funasr/tokenizer/abs_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..e125d292bdae20fe02af3f2e1ceddb2cd033f56b --- /dev/null +++ b/funasr/tokenizer/abs_tokenizer.py @@ -0,0 +1,93 @@ +import json +import numpy as np +from abc import ABC +from pathlib import Path +from abc import abstractmethod +from typing import Union, Iterable, List, Dict + + +class AbsTokenizer(ABC): + @abstractmethod + def text2tokens(self, line: str) -> List[str]: + raise NotImplementedError + + @abstractmethod + def tokens2text(self, tokens: Iterable[str]) -> str: + raise NotImplementedError + + +class BaseTokenizer(ABC): + def __init__( + self, + token_list: Union[Path, str, Iterable[str]] = None, + unk_symbol: str = "", + **kwargs, + ): + + if token_list is not None: + if isinstance(token_list, (Path, str)) and token_list.endswith(".txt"): + token_list = Path(token_list) + self.token_list_repr = str(token_list) + self.token_list: List[str] = [] + + with token_list.open("r", encoding="utf-8") as f: + for idx, line in enumerate(f): + line = line.rstrip() + self.token_list.append(line) + elif isinstance(token_list, (Path, str)) and token_list.endswith(".json"): + token_list = Path(token_list) + self.token_list_repr = str(token_list) + self.token_list: List[str] = [] + + with open(token_list, "r", encoding="utf-8") as f: + self.token_list = json.load(f) + + else: + self.token_list: List[str] = list(token_list) + self.token_list_repr = "" + for i, t in enumerate(self.token_list): + if i == 3: + break + self.token_list_repr += f"{t}, " + self.token_list_repr += f"... (NVocab={(len(self.token_list))})" + + self.token2id: Dict[str, int] = {} + for i, t in enumerate(self.token_list): + if t in self.token2id: + raise RuntimeError(f'Symbol "{t}" is duplicated') + self.token2id[t] = i + + self.unk_symbol = unk_symbol + if self.unk_symbol not in self.token2id: + raise RuntimeError(f"Unknown symbol '{unk_symbol}' doesn't exist in the token_list") + self.unk_id = self.token2id[self.unk_symbol] + + def encode(self, text, **kwargs): + tokens = self.text2tokens(text) + text_ints = self.tokens2ids(tokens) + + return text_ints + + def decode(self, text_ints): + token = self.ids2tokens(text_ints) + text = self.tokens2text(token) + return text + + def get_num_vocabulary_size(self) -> int: + return len(self.token_list) + + def ids2tokens(self, integers: Union[np.ndarray, Iterable[int]]) -> List[str]: + if isinstance(integers, np.ndarray) and integers.ndim != 1: + raise ValueError(f"Must be 1 dim ndarray, but got {integers.ndim}") + return [self.token_list[i] for i in integers] + + def tokens2ids(self, tokens: Iterable[str]) -> List[int]: + return [self.token2id.get(i, self.unk_id) for i in tokens] + + @abstractmethod + def text2tokens(self, line: str) -> List[str]: + raise NotImplementedError + + @abstractmethod + def tokens2text(self, tokens: Iterable[str]) -> str: + raise NotImplementedError diff --git a/funasr/tokenizer/build_tokenizer.py b/funasr/tokenizer/build_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..d6080f75d827604a51dadfcedda7aa493509dd42 --- /dev/null +++ b/funasr/tokenizer/build_tokenizer.py @@ -0,0 +1,59 @@ +from pathlib import Path +from typing import Iterable +from typing import Union + + +from funasr.tokenizer.abs_tokenizer import AbsTokenizer +from funasr.tokenizer.char_tokenizer import CharTokenizer +from funasr.tokenizer.phoneme_tokenizer import PhonemeTokenizer +from funasr.tokenizer.sentencepiece_tokenizer import SentencepiecesTokenizer +from funasr.tokenizer.word_tokenizer import WordTokenizer + + +def build_tokenizer( + token_type: str, + bpemodel: Union[Path, str, Iterable[str]] = None, + non_linguistic_symbols: Union[Path, str, Iterable[str]] = None, + remove_non_linguistic_symbols: bool = False, + space_symbol: str = "", + delimiter: str = None, + g2p_type: str = None, +) -> AbsTokenizer: + """A helper function to instantiate Tokenizer""" + if token_type == "bpe": + if bpemodel is None: + raise ValueError('bpemodel is required if token_type = "bpe"') + + if remove_non_linguistic_symbols: + raise RuntimeError( + "remove_non_linguistic_symbols is not implemented for token_type=bpe" + ) + return SentencepiecesTokenizer(bpemodel) + + elif token_type == "word": + if remove_non_linguistic_symbols and non_linguistic_symbols is not None: + return WordTokenizer( + delimiter=delimiter, + non_linguistic_symbols=non_linguistic_symbols, + remove_non_linguistic_symbols=True, + ) + else: + return WordTokenizer(delimiter=delimiter) + + elif token_type == "char": + return CharTokenizer( + non_linguistic_symbols=non_linguistic_symbols, + space_symbol=space_symbol, + remove_non_linguistic_symbols=remove_non_linguistic_symbols, + ) + + elif token_type == "phn": + return PhonemeTokenizer( + g2p_type=g2p_type, + non_linguistic_symbols=non_linguistic_symbols, + space_symbol=space_symbol, + remove_non_linguistic_symbols=remove_non_linguistic_symbols, + ) + + else: + raise ValueError(f"token_mode must be one of bpe, word, char or phn: " f"{token_type}") diff --git a/funasr/tokenizer/char_tokenizer.py b/funasr/tokenizer/char_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..7b517dacb5ca943cdff2676441baa4bc0062b8ba --- /dev/null +++ b/funasr/tokenizer/char_tokenizer.py @@ -0,0 +1,111 @@ +from pathlib import Path +from typing import Iterable +from typing import List +from typing import Union +import warnings +import re + +from funasr.tokenizer.abs_tokenizer import BaseTokenizer +from funasr.register import tables + + +@tables.register("tokenizer_classes", "CharTokenizer") +class CharTokenizer(BaseTokenizer): + def __init__( + self, + non_linguistic_symbols: Union[Path, str, Iterable[str]] = None, + space_symbol: str = "", + remove_non_linguistic_symbols: bool = False, + split_with_space: bool = False, + seg_dict: str = None, + **kwargs, + ): + super().__init__(**kwargs) + self.space_symbol = space_symbol + if non_linguistic_symbols is None: + self.non_linguistic_symbols = set() + elif isinstance(non_linguistic_symbols, (Path, str)): + non_linguistic_symbols = Path(non_linguistic_symbols) + try: + with non_linguistic_symbols.open("r", encoding="utf-8") as f: + self.non_linguistic_symbols = set(line.rstrip() for line in f) + except FileNotFoundError: + warnings.warn(f"{non_linguistic_symbols} doesn't exist.") + self.non_linguistic_symbols = set() + else: + self.non_linguistic_symbols = set(non_linguistic_symbols) + self.remove_non_linguistic_symbols = remove_non_linguistic_symbols + self.split_with_space = split_with_space + self.seg_dict = None + seg_dict = seg_dict if seg_dict is not None else kwargs.get("seg_dict_file", None) + if seg_dict is not None: + self.seg_dict = load_seg_dict(seg_dict) + + def __repr__(self): + return ( + f"{self.__class__.__name__}(" + f'space_symbol="{self.space_symbol}"' + f'non_linguistic_symbols="{self.non_linguistic_symbols}"' + f")" + ) + + def text2tokens(self, line: Union[str, list]) -> List[str]: + # if self.split_with_space: + if self.seg_dict is not None: + tokens = line.strip().split(" ") + tokens = seg_tokenize(tokens, self.seg_dict) + else: + tokens = [] + while len(line) != 0: + for w in self.non_linguistic_symbols: + if line.startswith(w): + if not self.remove_non_linguistic_symbols: + tokens.append(line[: len(w)]) + line = line[len(w) :] + break + else: + t = line[0] + if t == " ": + # t = "" + line = line[1:] + continue + tokens.append(t) + line = line[1:] + return tokens + + def tokens2text(self, tokens: Iterable[str]) -> str: + tokens = [t if t != self.space_symbol else " " for t in tokens] + return "".join(tokens) + + +def load_seg_dict(seg_dict_file): + seg_dict = {} + assert isinstance(seg_dict_file, str) + with open(seg_dict_file, "r", encoding="utf8") as f: + lines = f.readlines() + for line in lines: + s = line.strip().split() + key = s[0] + value = s[1:] + seg_dict[key] = " ".join(value) + return seg_dict + + +def seg_tokenize(txt, seg_dict): + # pattern = re.compile(r'^[\u4E00-\u9FA50-9]+$') + pattern = re.compile(r"([\u4E00-\u9FA5A-Za-z0-9])") + out_txt = "" + for word in txt: + word = word.lower() + if word in seg_dict: + out_txt += seg_dict[word] + " " + else: + if pattern.match(word): + for char in word: + if char in seg_dict: + out_txt += seg_dict[char] + " " + else: + out_txt += "" + " " + else: + out_txt += "" + " " + return out_txt.strip().split() diff --git a/funasr/tokenizer/cleaner.py b/funasr/tokenizer/cleaner.py new file mode 100644 index 0000000000000000000000000000000000000000..882db5c970c40350547a75d706f547ff12799430 --- /dev/null +++ b/funasr/tokenizer/cleaner.py @@ -0,0 +1,48 @@ +from typing import Collection + +from jaconv import jaconv + +# import tacotron_cleaner.cleaners + +try: + from vietnamese_cleaner import vietnamese_cleaners +except ImportError: + vietnamese_cleaners = None + + +class TextCleaner: + """Text cleaner. + + Examples: + >>> cleaner = TextCleaner("tacotron") + >>> cleaner("(Hello-World); & jr. & dr.") + 'HELLO WORLD, AND JUNIOR AND DOCTOR' + + """ + + def __init__(self, cleaner_types: Collection[str] = None): + + if cleaner_types is None: + self.cleaner_types = [] + elif isinstance(cleaner_types, str): + self.cleaner_types = [cleaner_types] + else: + self.cleaner_types = list(cleaner_types) + + def __call__(self, text: str) -> str: + for t in self.cleaner_types: + if t == "tacotron": + # text = tacotron_cleaner.cleaners.custom_english_cleaners(text) + pass + elif t == "jaconv": + text = jaconv.normalize(text) + elif t == "vietnamese": + if vietnamese_cleaners is None: + raise RuntimeError("Please install underthesea") + text = vietnamese_cleaners.vietnamese_cleaner(text) + elif t == "korean_cleaner": + text = KoreanCleaner.normalize_text(text) + else: + raise RuntimeError(f"Not supported: type={t}") + + return text diff --git a/funasr/tokenizer/hf_tokenizer.py b/funasr/tokenizer/hf_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..1ac7ae25988d165a24a627789f4997b247e31429 --- /dev/null +++ b/funasr/tokenizer/hf_tokenizer.py @@ -0,0 +1,13 @@ +from funasr.register import tables + + +@tables.register("tokenizer_classes", "HuggingfaceTokenizer") +def HuggingfaceTokenizer(init_param_path, **kwargs): + try: + from transformers import AutoTokenizer + except: + # print("If you want to use hugging, please `pip install -U transformers`") + pass + tokenizer = AutoTokenizer.from_pretrained(init_param_path) + + return tokenizer diff --git a/funasr/tokenizer/korean_cleaner.py b/funasr/tokenizer/korean_cleaner.py new file mode 100644 index 0000000000000000000000000000000000000000..eceebe21675be5a6738992fee506e66458449b3b --- /dev/null +++ b/funasr/tokenizer/korean_cleaner.py @@ -0,0 +1,74 @@ +# Referenced from https://github.com/hccho2/Tacotron-Wavenet-Vocoder-Korean + +import re + + +class KoreanCleaner: + @classmethod + def _normalize_numbers(cls, text): + number_to_kor = { + "0": "영", + "1": "일", + "2": "이", + "3": "삼", + "4": "사", + "5": "오", + "6": "육", + "7": "칠", + "8": "팔", + "9": "구", + } + new_text = "".join( + number_to_kor[char] if char in number_to_kor.keys() else char for char in text + ) + return new_text + + @classmethod + def _normalize_english_text(cls, text): + upper_alphabet_to_kor = { + "A": "에이", + "B": "비", + "C": "씨", + "D": "디", + "E": "이", + "F": "에프", + "G": "지", + "H": "에이치", + "I": "아이", + "J": "제이", + "K": "케이", + "L": "엘", + "M": "엠", + "N": "엔", + "O": "오", + "P": "피", + "Q": "큐", + "R": "알", + "S": "에스", + "T": "티", + "U": "유", + "V": "브이", + "W": "더블유", + "X": "엑스", + "Y": "와이", + "Z": "지", + } + new_text = re.sub("[a-z]+", lambda x: str.upper(x.group()), text) + new_text = "".join( + upper_alphabet_to_kor[char] if char in upper_alphabet_to_kor.keys() else char + for char in new_text + ) + + return new_text + + @classmethod + def normalize_text(cls, text): + # stage 0 : text strip + text = text.strip() + + # stage 1 : normalize numbers + text = cls._normalize_numbers(text) + + # stage 2 : normalize english text + text = cls._normalize_english_text(text) + return text diff --git a/funasr/tokenizer/phoneme_tokenizer.py b/funasr/tokenizer/phoneme_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..12b29ec220e40c73f3f26693a6647bb9a600fa82 --- /dev/null +++ b/funasr/tokenizer/phoneme_tokenizer.py @@ -0,0 +1,524 @@ +import logging +from pathlib import Path +import re +from typing import Iterable +from typing import List +from typing import Optional +from typing import Union +import warnings + +# import g2p_en +import jamo + +from funasr.tokenizer.abs_tokenizer import AbsTokenizer + + +g2p_classes = [ + None, + "g2p_en", + "g2p_en_no_space", + "pyopenjtalk", + "pyopenjtalk_kana", + "pyopenjtalk_accent", + "pyopenjtalk_accent_with_pause", + "pyopenjtalk_prosody", + "pypinyin_g2p", + "pypinyin_g2p_phone", + "espeak_ng_arabic", + "espeak_ng_german", + "espeak_ng_french", + "espeak_ng_spanish", + "espeak_ng_russian", + "espeak_ng_greek", + "espeak_ng_finnish", + "espeak_ng_hungarian", + "espeak_ng_dutch", + "espeak_ng_english_us_vits", + "espeak_ng_hindi", + "g2pk", + "g2pk_no_space", + "korean_jaso", + "korean_jaso_no_space", +] + + +def split_by_space(text) -> List[str]: + if " " in text: + text = text.replace(" ", " ") + return [c.replace("", " ") for c in text.split(" ")] + else: + return text.split(" ") + + +def pyopenjtalk_g2p(text) -> List[str]: + import pyopenjtalk + + # phones is a str object separated by space + phones = pyopenjtalk.g2p(text, kana=False) + phones = phones.split(" ") + return phones + + +def pyopenjtalk_g2p_accent(text) -> List[str]: + import pyopenjtalk + import re + + phones = [] + for labels in pyopenjtalk.run_frontend(text)[1]: + p = re.findall(r"\-(.*?)\+.*?\/A:([0-9\-]+).*?\/F:.*?_([0-9]+)", labels) + if len(p) == 1: + phones += [p[0][0], p[0][2], p[0][1]] + return phones + + +def pyopenjtalk_g2p_accent_with_pause(text) -> List[str]: + import pyopenjtalk + import re + + phones = [] + for labels in pyopenjtalk.run_frontend(text)[1]: + if labels.split("-")[1].split("+")[0] == "pau": + phones += ["pau"] + continue + p = re.findall(r"\-(.*?)\+.*?\/A:([0-9\-]+).*?\/F:.*?_([0-9]+)", labels) + if len(p) == 1: + phones += [p[0][0], p[0][2], p[0][1]] + return phones + + +def pyopenjtalk_g2p_kana(text) -> List[str]: + import pyopenjtalk + + kanas = pyopenjtalk.g2p(text, kana=True) + return list(kanas) + + +def pyopenjtalk_g2p_prosody(text: str, drop_unvoiced_vowels: bool = True) -> List[str]: + """Extract phoneme + prosoody symbol sequence from input full-context labels. + + The algorithm is based on `Prosodic features control by symbols as input of + sequence-to-sequence acoustic modeling for neural TTS`_ with some r9y9's tweaks. + + Args: + text (str): Input text. + drop_unvoiced_vowels (bool): whether to drop unvoiced vowels. + + Returns: + List[str]: List of phoneme + prosody symbols. + + Examples: + >>> from funasr.tokenizer.phoneme_tokenizer import pyopenjtalk_g2p_prosody + >>> pyopenjtalk_g2p_prosody("こんにちは。") + ['^', 'k', 'o', '[', 'N', 'n', 'i', 'ch', 'i', 'w', 'a', '$'] + + .. _`Prosodic features control by symbols as input of sequence-to-sequence acoustic + modeling for neural TTS`: https://doi.org/10.1587/transinf.2020EDP7104 + + """ + import pyopenjtalk + + labels = pyopenjtalk.run_frontend(text)[1] + N = len(labels) + + phones = [] + for n in range(N): + lab_curr = labels[n] + + # current phoneme + p3 = re.search(r"\-(.*?)\+", lab_curr).group(1) + + # deal unvoiced vowels as normal vowels + if drop_unvoiced_vowels and p3 in "AEIOU": + p3 = p3.lower() + + # deal with sil at the beginning and the end of text + if p3 == "sil": + assert n == 0 or n == N - 1 + if n == 0: + phones.append("^") + elif n == N - 1: + # check question form or not + e3 = _numeric_feature_by_regex(r"!(\d+)_", lab_curr) + if e3 == 0: + phones.append("$") + elif e3 == 1: + phones.append("?") + continue + elif p3 == "pau": + phones.append("_") + continue + else: + phones.append(p3) + + # accent type and position info (forward or backward) + a1 = _numeric_feature_by_regex(r"/A:([0-9\-]+)\+", lab_curr) + a2 = _numeric_feature_by_regex(r"\+(\d+)\+", lab_curr) + a3 = _numeric_feature_by_regex(r"\+(\d+)/", lab_curr) + + # number of mora in accent phrase + f1 = _numeric_feature_by_regex(r"/F:(\d+)_", lab_curr) + + a2_next = _numeric_feature_by_regex(r"\+(\d+)\+", labels[n + 1]) + # accent phrase border + if a3 == 1 and a2_next == 1 and p3 in "aeiouAEIOUNcl": + phones.append("#") + # pitch falling + elif a1 == 0 and a2_next == a2 + 1 and a2 != f1: + phones.append("]") + # pitch rising + elif a2 == 1 and a2_next == 2: + phones.append("[") + + return phones + + +def _numeric_feature_by_regex(regex, s): + match = re.search(regex, s) + if match is None: + return -50 + return int(match.group(1)) + + +def pypinyin_g2p(text) -> List[str]: + from pypinyin import pinyin + from pypinyin import Style + + phones = [phone[0] for phone in pinyin(text, style=Style.TONE3)] + return phones + + +def pypinyin_g2p_phone(text) -> List[str]: + from pypinyin import pinyin + from pypinyin import Style + from pypinyin.style._utils import get_finals + from pypinyin.style._utils import get_initials + + phones = [ + p + for phone in pinyin(text, style=Style.TONE3) + for p in [ + get_initials(phone[0], strict=True), + get_finals(phone[0], strict=True), + ] + if len(p) != 0 + ] + return phones + + +class G2p_en: + """On behalf of g2p_en.G2p. + + g2p_en.G2p isn't pickalable and it can't be copied to the other processes + via multiprocessing module. + As a workaround, g2p_en.G2p is instantiated upon calling this class. + + """ + + def __init__(self, no_space: bool = False): + self.no_space = no_space + self.g2p = None + + def __call__(self, text) -> List[str]: + if self.g2p is None: + self.g2p = g2p_en.G2p() + + phones = self.g2p(text) + if self.no_space: + # remove space which represents word serapater + phones = list(filter(lambda s: s != " ", phones)) + return phones + + +class G2pk: + """On behalf of g2pk.G2p. + + g2pk.G2p isn't pickalable and it can't be copied to the other processes + via multiprocessing module. + As a workaround, g2pk.G2p is instantiated upon calling this class. + + """ + + def __init__(self, descritive=False, group_vowels=False, to_syl=False, no_space=False): + self.descritive = descritive + self.group_vowels = group_vowels + self.to_syl = to_syl + self.no_space = no_space + self.g2p = None + + def __call__(self, text) -> List[str]: + if self.g2p is None: + import g2pk + + self.g2p = g2pk.G2p() + + phones = list( + self.g2p( + text, + descriptive=self.descritive, + group_vowels=self.group_vowels, + to_syl=self.to_syl, + ) + ) + if self.no_space: + # remove space which represents word serapater + phones = list(filter(lambda s: s != " ", phones)) + return phones + + +class Jaso: + PUNC = "!'(),-.:;?" + SPACE = " " + + JAMO_LEADS = "".join([chr(_) for _ in range(0x1100, 0x1113)]) + JAMO_VOWELS = "".join([chr(_) for _ in range(0x1161, 0x1176)]) + JAMO_TAILS = "".join([chr(_) for _ in range(0x11A8, 0x11C3)]) + + VALID_CHARS = JAMO_LEADS + JAMO_VOWELS + JAMO_TAILS + PUNC + SPACE + + def __init__(self, space_symbol=" ", no_space=False): + self.space_symbol = space_symbol + self.no_space = no_space + + def _text_to_jaso(self, line: str) -> List[str]: + jasos = list(jamo.hangul_to_jamo(line)) + return jasos + + def _remove_non_korean_characters(self, tokens): + new_tokens = [token for token in tokens if token in self.VALID_CHARS] + return new_tokens + + def __call__(self, text) -> List[str]: + graphemes = [x for x in self._text_to_jaso(text)] + graphemes = self._remove_non_korean_characters(graphemes) + + if self.no_space: + graphemes = list(filter(lambda s: s != " ", graphemes)) + else: + graphemes = [x if x != " " else self.space_symbol for x in graphemes] + return graphemes + + +class Phonemizer: + """Phonemizer module for various languages. + + This is wrapper module of https://github.com/bootphon/phonemizer. + You can define various g2p modules by specifying options for phonemizer. + + See available options: + https://github.com/bootphon/phonemizer/blob/master/phonemizer/phonemize.py#L32 + + """ + + def __init__( + self, + backend, + word_separator: Optional[str] = None, + syllable_separator: Optional[str] = None, + phone_separator: Optional[str] = " ", + strip=False, + split_by_single_token: bool = False, + **phonemizer_kwargs, + ): + # delayed import + from phonemizer.backend import BACKENDS + from phonemizer.separator import Separator + + self.separator = Separator( + word=word_separator, + syllable=syllable_separator, + phone=phone_separator, + ) + + # define logger to suppress the warning in phonemizer + logger = logging.getLogger("phonemizer") + logger.setLevel(logging.ERROR) + self.phonemizer = BACKENDS[backend]( + **phonemizer_kwargs, + logger=logger, + ) + self.strip = strip + self.split_by_single_token = split_by_single_token + + def __call__(self, text) -> List[str]: + tokens = self.phonemizer.phonemize( + [text], + separator=self.separator, + strip=self.strip, + njobs=1, + )[0] + if not self.split_by_single_token: + return tokens.split() + else: + # "a: ab" -> ["a", ":", "", "a", "b"] + # TODO(kan-bayashi): space replacement should be dealt in PhonemeTokenizer + return [c.replace(" ", "") for c in tokens] + + +class PhonemeTokenizer(AbsTokenizer): + def __init__( + self, + g2p_type: Union[None, str], + non_linguistic_symbols: Union[Path, str, Iterable[str]] = None, + space_symbol: str = "", + remove_non_linguistic_symbols: bool = False, + ): + if g2p_type is None: + self.g2p = split_by_space + elif g2p_type == "g2p_en": + self.g2p = G2p_en(no_space=False) + elif g2p_type == "g2p_en_no_space": + self.g2p = G2p_en(no_space=True) + elif g2p_type == "pyopenjtalk": + self.g2p = pyopenjtalk_g2p + elif g2p_type == "pyopenjtalk_kana": + self.g2p = pyopenjtalk_g2p_kana + elif g2p_type == "pyopenjtalk_accent": + self.g2p = pyopenjtalk_g2p_accent + elif g2p_type == "pyopenjtalk_accent_with_pause": + self.g2p = pyopenjtalk_g2p_accent_with_pause + elif g2p_type == "pyopenjtalk_prosody": + self.g2p = pyopenjtalk_g2p_prosody + elif g2p_type == "pypinyin_g2p": + self.g2p = pypinyin_g2p + elif g2p_type == "pypinyin_g2p_phone": + self.g2p = pypinyin_g2p_phone + elif g2p_type == "espeak_ng_arabic": + self.g2p = Phonemizer( + language="ar", + backend="espeak", + with_stress=True, + preserve_punctuation=True, + ) + elif g2p_type == "espeak_ng_german": + self.g2p = Phonemizer( + language="de", + backend="espeak", + with_stress=True, + preserve_punctuation=True, + ) + elif g2p_type == "espeak_ng_french": + self.g2p = Phonemizer( + language="fr-fr", + backend="espeak", + with_stress=True, + preserve_punctuation=True, + ) + elif g2p_type == "espeak_ng_spanish": + self.g2p = Phonemizer( + language="es", + backend="espeak", + with_stress=True, + preserve_punctuation=True, + ) + elif g2p_type == "espeak_ng_russian": + self.g2p = Phonemizer( + language="ru", + backend="espeak", + with_stress=True, + preserve_punctuation=True, + ) + elif g2p_type == "espeak_ng_greek": + self.g2p = Phonemizer( + language="el", + backend="espeak", + with_stress=True, + preserve_punctuation=True, + ) + elif g2p_type == "espeak_ng_finnish": + self.g2p = Phonemizer( + language="fi", + backend="espeak", + with_stress=True, + preserve_punctuation=True, + ) + elif g2p_type == "espeak_ng_hungarian": + self.g2p = Phonemizer( + language="hu", + backend="espeak", + with_stress=True, + preserve_punctuation=True, + ) + elif g2p_type == "espeak_ng_dutch": + self.g2p = Phonemizer( + language="nl", + backend="espeak", + with_stress=True, + preserve_punctuation=True, + ) + elif g2p_type == "espeak_ng_hindi": + self.g2p = Phonemizer( + language="hi", + backend="espeak", + with_stress=True, + preserve_punctuation=True, + ) + elif g2p_type == "g2pk": + self.g2p = G2pk(no_space=False) + elif g2p_type == "g2pk_no_space": + self.g2p = G2pk(no_space=True) + elif g2p_type == "espeak_ng_english_us_vits": + # VITS official implementation-like processing + # Reference: https://github.com/jaywalnut310/vits + self.g2p = Phonemizer( + language="en-us", + backend="espeak", + with_stress=True, + preserve_punctuation=True, + strip=True, + word_separator=" ", + phone_separator="", + split_by_single_token=True, + ) + elif g2p_type == "korean_jaso": + self.g2p = Jaso(space_symbol=space_symbol, no_space=False) + elif g2p_type == "korean_jaso_no_space": + self.g2p = Jaso(no_space=True) + else: + raise NotImplementedError(f"Not supported: g2p_type={g2p_type}") + + self.g2p_type = g2p_type + self.space_symbol = space_symbol + if non_linguistic_symbols is None: + self.non_linguistic_symbols = set() + elif isinstance(non_linguistic_symbols, (Path, str)): + non_linguistic_symbols = Path(non_linguistic_symbols) + try: + with non_linguistic_symbols.open("r", encoding="utf-8") as f: + self.non_linguistic_symbols = set(line.rstrip() for line in f) + except FileNotFoundError: + warnings.warn(f"{non_linguistic_symbols} doesn't exist.") + self.non_linguistic_symbols = set() + else: + self.non_linguistic_symbols = set(non_linguistic_symbols) + self.remove_non_linguistic_symbols = remove_non_linguistic_symbols + + def __repr__(self): + return ( + f"{self.__class__.__name__}(" + f'g2p_type="{self.g2p_type}", ' + f'space_symbol="{self.space_symbol}", ' + f'non_linguistic_symbols="{self.non_linguistic_symbols}"' + ")" + ) + + def text2tokens(self, line: str) -> List[str]: + tokens = [] + while len(line) != 0: + for w in self.non_linguistic_symbols: + if line.startswith(w): + if not self.remove_non_linguistic_symbols: + tokens.append(line[: len(w)]) + line = line[len(w) :] + break + else: + t = line[0] + tokens.append(t) + line = line[1:] + + line = "".join(tokens) + tokens = self.g2p(line) + return tokens + + def tokens2text(self, tokens: Iterable[str]) -> str: + # phoneme type is not invertible + return "".join(tokens) diff --git a/funasr/tokenizer/sentencepiece_tokenizer.py b/funasr/tokenizer/sentencepiece_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..0b47a9f6afc2b50b6ea7d2990510df461c71dd5e --- /dev/null +++ b/funasr/tokenizer/sentencepiece_tokenizer.py @@ -0,0 +1,57 @@ +from pathlib import Path +from typing import Iterable +from typing import List +from typing import Union + +import sentencepiece as spm + +from funasr.tokenizer.abs_tokenizer import BaseTokenizer +from funasr.register import tables + + +@tables.register("tokenizer_classes", "SentencepiecesTokenizer") +class SentencepiecesTokenizer(BaseTokenizer): + def __init__(self, bpemodel: Union[Path, str], **kwargs): + super().__init__(**kwargs) + self.bpemodel = str(bpemodel) + # NOTE(kamo): + # Don't build SentencePieceProcessor in __init__() + # because it's not picklable and it may cause following error, + # "TypeError: can't pickle SwigPyObject objects", + # when giving it as argument of "multiprocessing.Process()". + self.sp = None + self._build_sentence_piece_processor() + + def __repr__(self): + return f'{self.__class__.__name__}(model="{self.bpemodel}")' + + def _build_sentence_piece_processor(self): + # Build SentencePieceProcessor lazily. + if self.sp is None: + self.sp = spm.SentencePieceProcessor() + self.sp.load(self.bpemodel) + + def text2tokens(self, line: str) -> List[str]: + self._build_sentence_piece_processor() + return self.sp.EncodeAsPieces(line) + + def tokens2text(self, tokens: Iterable[str]) -> str: + self._build_sentence_piece_processor() + return self.sp.DecodePieces(list(tokens)) + + def encode(self, line: str, **kwargs) -> List[int]: + self._build_sentence_piece_processor() + return self.sp.EncodeAsIds(line) + + def decode(self, line: List[int], **kwargs): + self._build_sentence_piece_processor() + return self.sp.DecodeIds(line) + + def get_vocab_size(self): + return self.sp.GetPieceSize() + + def ids2tokens(self, *args, **kwargs): + return self.decode(*args, **kwargs) + + def tokens2ids(self, *args, **kwargs): + return self.encode(*args, **kwargs) diff --git a/funasr/tokenizer/token_id_converter.py b/funasr/tokenizer/token_id_converter.py new file mode 100644 index 0000000000000000000000000000000000000000..c3fc71d5a363192683b108b113fec3ab7799bdc4 --- /dev/null +++ b/funasr/tokenizer/token_id_converter.py @@ -0,0 +1,56 @@ +from pathlib import Path +from typing import Dict +from typing import Iterable +from typing import List +from typing import Union + +import numpy as np + + +class TokenIDConverter: + def __init__( + self, + token_list: Union[Path, str, Iterable[str]], + unk_symbol: str = "", + ): + + if isinstance(token_list, (Path, str)): + token_list = Path(token_list) + self.token_list_repr = str(token_list) + self.token_list: List[str] = [] + + with token_list.open("r", encoding="utf-8") as f: + for idx, line in enumerate(f): + line = line.rstrip() + self.token_list.append(line) + + else: + self.token_list: List[str] = list(token_list) + self.token_list_repr = "" + for i, t in enumerate(self.token_list): + if i == 3: + break + self.token_list_repr += f"{t}, " + self.token_list_repr += f"... (NVocab={(len(self.token_list))})" + + self.token2id: Dict[str, int] = {} + for i, t in enumerate(self.token_list): + if t in self.token2id: + raise RuntimeError(f'Symbol "{t}" is duplicated') + self.token2id[t] = i + + self.unk_symbol = unk_symbol + if self.unk_symbol not in self.token2id: + raise RuntimeError(f"Unknown symbol '{unk_symbol}' doesn't exist in the token_list") + self.unk_id = self.token2id[self.unk_symbol] + + def get_num_vocabulary_size(self) -> int: + return len(self.token_list) + + def ids2tokens(self, integers: Union[np.ndarray, Iterable[int]]) -> List[str]: + if isinstance(integers, np.ndarray) and integers.ndim != 1: + raise ValueError(f"Must be 1 dim ndarray, but got {integers.ndim}") + return [self.token_list[i] for i in integers] + + def tokens2ids(self, tokens: Iterable[str]) -> List[int]: + return [self.token2id.get(i, self.unk_id) for i in tokens] diff --git a/funasr/tokenizer/whisper_tokenizer.py b/funasr/tokenizer/whisper_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..bd96b832b9b6c2bce9b53232a5240dc977a82541 --- /dev/null +++ b/funasr/tokenizer/whisper_tokenizer.py @@ -0,0 +1,45 @@ +from funasr.register import tables + + +@tables.register("tokenizer_classes", "WhisperTokenizer") +def WhisperTokenizer(**kwargs): + try: + from whisper.tokenizer import get_tokenizer + except: + print("Notice: If you want to use whisper, please `pip install -U openai-whisper`") + + language = kwargs.get("language", None) + task = kwargs.get("task", "transcribe") + is_multilingual = kwargs.get("is_multilingual", True) + num_languages = kwargs.get("num_languages", 99) + tokenizer = get_tokenizer( + multilingual=is_multilingual, + num_languages=num_languages, + language=language, + task=task, + ) + + return tokenizer + + +@tables.register("tokenizer_classes", "SenseVoiceTokenizer") +def SenseVoiceTokenizer(**kwargs): + try: + from funasr.models.sense_voice.whisper_lib.tokenizer import get_tokenizer + except: + print("Notice: If you want to use whisper, please `pip install -U openai-whisper`") + + language = kwargs.get("language", None) + task = kwargs.get("task", None) + is_multilingual = kwargs.get("is_multilingual", True) + num_languages = kwargs.get("num_languages", 8749) + vocab_path = kwargs.get("vocab_path", None) + tokenizer = get_tokenizer( + multilingual=is_multilingual, + num_languages=num_languages, + language=language, + task=task, + vocab_path=vocab_path, + ) + + return tokenizer diff --git a/funasr/tokenizer/word_tokenizer.py b/funasr/tokenizer/word_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..c1fd6e039d3f5b47a986c1e47d2d8794b3dcca53 --- /dev/null +++ b/funasr/tokenizer/word_tokenizer.py @@ -0,0 +1,55 @@ +from pathlib import Path +from typing import Iterable +from typing import List +from typing import Union +import warnings + + +from funasr.tokenizer.abs_tokenizer import AbsTokenizer + + +class WordTokenizer(AbsTokenizer): + def __init__( + self, + delimiter: str = None, + non_linguistic_symbols: Union[Path, str, Iterable[str]] = None, + remove_non_linguistic_symbols: bool = False, + ): + self.delimiter = delimiter + + if not remove_non_linguistic_symbols and non_linguistic_symbols is not None: + warnings.warn( + "non_linguistic_symbols is only used " "when remove_non_linguistic_symbols = True" + ) + + if non_linguistic_symbols is None: + self.non_linguistic_symbols = set() + elif isinstance(non_linguistic_symbols, (Path, str)): + non_linguistic_symbols = Path(non_linguistic_symbols) + try: + with non_linguistic_symbols.open("r", encoding="utf-8") as f: + self.non_linguistic_symbols = set(line.rstrip() for line in f) + except FileNotFoundError: + warnings.warn(f"{non_linguistic_symbols} doesn't exist.") + self.non_linguistic_symbols = set() + else: + self.non_linguistic_symbols = set(non_linguistic_symbols) + self.remove_non_linguistic_symbols = remove_non_linguistic_symbols + + def __repr__(self): + return f'{self.__class__.__name__}(delimiter="{self.delimiter}")' + + def text2tokens(self, line: str) -> List[str]: + tokens = [] + for t in line.split(self.delimiter): + if self.remove_non_linguistic_symbols and t in self.non_linguistic_symbols: + continue + tokens.append(t) + return tokens + + def tokens2text(self, tokens: Iterable[str]) -> str: + if self.delimiter is None: + delimiter = " " + else: + delimiter = self.delimiter + return delimiter.join(tokens) diff --git a/funasr/train_utils/__init__.py b/funasr/train_utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/train_utils/add_gradient_noise.py b/funasr/train_utils/add_gradient_noise.py new file mode 100644 index 0000000000000000000000000000000000000000..654e928ec003b9592120451230c1e2253b2b09c0 --- /dev/null +++ b/funasr/train_utils/add_gradient_noise.py @@ -0,0 +1,31 @@ +import torch + + +def add_gradient_noise( + model: torch.nn.Module, + iteration: int, + duration: float = 100, + eta: float = 1.0, + scale_factor: float = 0.55, +): + """Adds noise from a standard normal distribution to the gradients. + + The standard deviation (`sigma`) is controlled + by the three hyper-parameters below. + `sigma` goes to zero (no noise) with more iterations. + + Args: + model: Model. + iteration: Number of iterations. + duration: {100, 1000}: Number of durations to control + the interval of the `sigma` change. + eta: {0.01, 0.3, 1.0}: The magnitude of `sigma`. + scale_factor: {0.55}: The scale of `sigma`. + """ + interval = (iteration // duration) + 1 + sigma = eta / interval**scale_factor + for param in model.parameters(): + if param.grad is not None: + _shape = param.grad.size() + noise = sigma * torch.randn(_shape).to(param.device) + param.grad += noise diff --git a/funasr/train_utils/average_nbest_models.py b/funasr/train_utils/average_nbest_models.py new file mode 100644 index 0000000000000000000000000000000000000000..873f4192c15a8eccc2f4f607de4493e68d718443 --- /dev/null +++ b/funasr/train_utils/average_nbest_models.py @@ -0,0 +1,97 @@ +import logging +from pathlib import Path +from typing import Optional +from typing import Sequence +from typing import Union +import warnings +import os +from io import BytesIO + +import torch +from typing import Collection +import os +import torch +import re +from collections import OrderedDict +from functools import cmp_to_key + + +def _get_checkpoint_paths(output_dir: str, last_n: int = 5, use_deepspeed=False, **kwargs): + """ + Get the paths of the last 'last_n' checkpoints by parsing filenames + in the output directory. + """ + try: + if not use_deepspeed: + checkpoint = torch.load(os.path.join(output_dir, "model.pt"), map_location="cpu") + else: + checkpoint = torch.load( + os.path.join(output_dir, "model.pt", "mp_rank_00_model_states.pt"), + map_location="cpu", + ) + avg_keep_nbest_models_type = checkpoint["avg_keep_nbest_models_type"] + val_step_or_epoch = checkpoint[f"val_{avg_keep_nbest_models_type}_step_or_epoch"] + sorted_items = sorted(val_step_or_epoch.items(), key=lambda x: x[1], reverse=True) + sorted_items = ( + sorted_items[:last_n] if avg_keep_nbest_models_type == "acc" else sorted_items[-last_n:] + ) + checkpoint_paths = [] + for key, value in sorted_items[:last_n]: + if not use_deepspeed: + ckpt = os.path.join(output_dir, key) + else: + ckpt = os.path.join(output_dir, key, "mp_rank_00_model_states.pt") + checkpoint_paths.append(ckpt) + + except: + print(f"{checkpoint} does not exist, avg the lastet checkpoint.") + # List all files in the output directory + files = os.listdir(output_dir) + # Filter out checkpoint files and extract epoch numbers + checkpoint_files = [f for f in files if f.startswith("model.pt.e")] + # Sort files by epoch number in descending order + checkpoint_files.sort(key=lambda x: int(re.search(r"(\d+)", x).group()), reverse=True) + # Get the last 'last_n' checkpoint paths + checkpoint_paths = [os.path.join(output_dir, f) for f in checkpoint_files[:last_n]] + + return checkpoint_paths + + +@torch.no_grad() +def average_checkpoints(output_dir: str, last_n: int = 5, **kwargs): + """ + Average the last 'last_n' checkpoints' model state_dicts. + If a tensor is of type torch.int, perform sum instead of average. + """ + checkpoint_paths = _get_checkpoint_paths(output_dir, last_n, **kwargs) + print(f"average_checkpoints: {checkpoint_paths}") + state_dicts = [] + + # Load state_dicts from checkpoints + for path in checkpoint_paths: + if os.path.isfile(path): + state_dicts.append(torch.load(path, map_location="cpu")["state_dict"]) + else: + print(f"Checkpoint file {path} not found.") + + # Check if we have any state_dicts to average + if len(state_dicts) < 1: + print("No checkpoints found for averaging.") + return + + # Average or sum weights + avg_state_dict = OrderedDict() + for key in state_dicts[0].keys(): + tensors = [state_dict[key].cpu() for state_dict in state_dicts] + # Check the type of the tensor + if str(tensors[0].dtype).startswith("torch.int"): + # Perform sum for integer tensors + summed_tensor = sum(tensors) + avg_state_dict[key] = summed_tensor + else: + # Perform average for other types of tensors + stacked_tensors = torch.stack(tensors) + avg_state_dict[key] = torch.mean(stacked_tensors, dim=0) + checkpoint_outpath = os.path.join(output_dir, f"model.pt.avg{last_n}") + torch.save({"state_dict": avg_state_dict}, checkpoint_outpath) + return checkpoint_outpath diff --git a/funasr/train_utils/device_funcs.py b/funasr/train_utils/device_funcs.py new file mode 100644 index 0000000000000000000000000000000000000000..dd7fc76a2862a273d90b0d0aaeb6ad88243a52ff --- /dev/null +++ b/funasr/train_utils/device_funcs.py @@ -0,0 +1,64 @@ +import dataclasses +import warnings + +import numpy as np +import torch + + +def to_device(data, device=None, dtype=None, non_blocking=False, copy=False): + """Change the device of object recursively""" + if isinstance(data, dict): + return {k: to_device(v, device, dtype, non_blocking, copy) for k, v in data.items()} + elif dataclasses.is_dataclass(data) and not isinstance(data, type): + return type(data)( + *[to_device(v, device, dtype, non_blocking, copy) for v in dataclasses.astuple(data)] + ) + # maybe namedtuple. I don't know the correct way to judge namedtuple. + elif isinstance(data, tuple) and type(data) is not tuple: + return type(data)(*[to_device(o, device, dtype, non_blocking, copy) for o in data]) + elif isinstance(data, (list, tuple)): + return type(data)(to_device(v, device, dtype, non_blocking, copy) for v in data) + elif isinstance(data, np.ndarray): + return to_device(torch.from_numpy(data), device, dtype, non_blocking, copy) + elif isinstance(data, torch.Tensor): + return data.to(device, dtype, non_blocking, copy) + else: + return data + + +def force_gatherable(data, device): + """Change object to gatherable in torch.nn.DataParallel recursively + + The difference from to_device() is changing to torch.Tensor if float or int + value is found. + + The restriction to the returned value in DataParallel: + The object must be + - torch.cuda.Tensor + - 1 or more dimension. 0-dimension-tensor sends warning. + or a list, tuple, dict. + + """ + if isinstance(data, dict): + return {k: force_gatherable(v, device) for k, v in data.items()} + # DataParallel can't handle NamedTuple well + elif isinstance(data, tuple) and type(data) is not tuple: + return type(data)(*[force_gatherable(o, device) for o in data]) + elif isinstance(data, (list, tuple, set)): + return type(data)(force_gatherable(v, device) for v in data) + elif isinstance(data, np.ndarray): + return force_gatherable(torch.from_numpy(data), device) + elif isinstance(data, torch.Tensor): + if data.dim() == 0: + # To 1-dim array + data = data[None] + return data.to(device) + elif isinstance(data, float): + return torch.tensor([data], dtype=torch.float, device=device) + elif isinstance(data, int): + return torch.tensor([data], dtype=torch.long, device=device) + elif data is None: + return None + else: + warnings.warn(f"{type(data)} may not be gatherable by DataParallel") + return data diff --git a/funasr/train_utils/forward_adaptor.py b/funasr/train_utils/forward_adaptor.py new file mode 100644 index 0000000000000000000000000000000000000000..eb6da2bb189690e85bc9192b5bd1958854a1353c --- /dev/null +++ b/funasr/train_utils/forward_adaptor.py @@ -0,0 +1,31 @@ +import torch + + +class ForwardAdaptor(torch.nn.Module): + """Wrapped module to parallelize specified method + + torch.nn.DataParallel parallelizes only "forward()" + and, maybe, the method having the other name can't be applied + except for wrapping the module just like this class. + + Examples: + >>> class A(torch.nn.Module): + ... def foo(self, x): + ... ... + >>> model = A() + >>> model = ForwardAdaptor(model, "foo") + >>> model = torch.nn.DataParallel(model, device_ids=[0, 1]) + >>> x = torch.randn(2, 10) + >>> model(x) + """ + + def __init__(self, module: torch.nn.Module, name: str): + super().__init__() + self.module = module + self.name = name + if not hasattr(module, name): + raise ValueError(f"{module} doesn't have {name}") + + def forward(self, *args, **kwargs): + func = getattr(self.module, self.name) + return func(*args, **kwargs) diff --git a/funasr/train_utils/initialize.py b/funasr/train_utils/initialize.py new file mode 100644 index 0000000000000000000000000000000000000000..667585a27412838823da5b63626097bdf1706f67 --- /dev/null +++ b/funasr/train_utils/initialize.py @@ -0,0 +1,55 @@ +#!/usr/bin/env python3 + +"""Initialize modules for espnet2 neural networks.""" + +import math +import torch + + +def initialize(model: torch.nn.Module, init: str): + """Initialize weights of a neural network module. + + Parameters are initialized using the given method or distribution. + + Custom initialization routines can be implemented into submodules + as function `espnet_initialization_fn` within the custom module. + + Args: + model: Target. + init: Method of initialization. + """ + + # weight init + for p in model.parameters(): + if p.dim() > 1: + if init == "xavier_uniform": + torch.nn.init.xavier_uniform_(p.data) + elif init == "xavier_normal": + torch.nn.init.xavier_normal_(p.data) + elif init == "kaiming_uniform": + torch.nn.init.kaiming_uniform_(p.data, nonlinearity="relu") + elif init == "kaiming_normal": + torch.nn.init.kaiming_normal_(p.data, nonlinearity="relu") + else: + raise ValueError("Unknown initialization: " + init) + # bias init + for p in model.parameters(): + if p.dim() == 1: + p.data.zero_() + + # reset some modules with default init + for m in model.modules(): + if isinstance(m, (torch.nn.Embedding, torch.nn.LayerNorm, torch.nn.GroupNorm)): + m.reset_parameters() + if hasattr(m, "espnet_initialization_fn"): + m.espnet_initialization_fn() + + # TODO(xkc): Hacking s3prl_frontend and wav2vec2encoder initialization + if getattr(model, "encoder", None) and getattr( + model.encoder, "reload_pretrained_parameters", None + ): + model.encoder.reload_pretrained_parameters() + if getattr(model, "frontend", None) and getattr( + model.frontend, "reload_pretrained_parameters", None + ): + model.frontend.reload_pretrained_parameters() diff --git a/funasr/train_utils/load_pretrained_model.py b/funasr/train_utils/load_pretrained_model.py new file mode 100644 index 0000000000000000000000000000000000000000..159ce2691f776d3714667ea073183919e3fc4c51 --- /dev/null +++ b/funasr/train_utils/load_pretrained_model.py @@ -0,0 +1,103 @@ +from typing import Any +from typing import Dict +from typing import Union +from io import BytesIO + +import logging +import torch +import torch.nn +import torch.optim +import pdb +import copy + + +def load_pretrained_model( + path: str, + model: torch.nn.Module, + ignore_init_mismatch: bool = True, + map_location: str = "cpu", + oss_bucket=None, + scope_map=[], + excludes=None, + **kwargs, +): + """Load a model state and set it to the model. + + Args: + init_param: ::: + + Examples: + + """ + + obj = model + dst_state = obj.state_dict() + + logging.info(f"ckpt: {path}") + + if oss_bucket is None: + ori_state = torch.load(path, map_location=map_location, weights_only=True) + else: + buffer = BytesIO(oss_bucket.get_object(path).read()) + ori_state = torch.load(buffer, map_location=map_location) + + src_state = copy.deepcopy(ori_state) + src_state = src_state["state_dict"] if "state_dict" in src_state else src_state + src_state = src_state["model_state_dict"] if "model_state_dict" in src_state else src_state + src_state = src_state["model"] if "model" in src_state else src_state + + if isinstance(scope_map, str): + scope_map = scope_map.split(",") + scope_map += ["module.", "None"] + logging.info(f"scope_map: {scope_map}") + + if excludes is not None: + if isinstance(excludes, str): + excludes = excludes.split(",") + + logging.info(f"excludes: {excludes}") + + for k in dst_state.keys(): + excludes_flag = False + if excludes is not None: + for k_ex in excludes: + if k.startswith(k_ex): + logging.info(f"key: {k} matching: {k_ex}, excluded") + excludes_flag = True + break + if excludes_flag: + continue + + k_src = k + + if scope_map is not None: + src_prefix = "" + dst_prefix = "" + for i in range(0, len(scope_map), 2): + src_prefix = scope_map[i] if scope_map[i].lower() != "none" else "" + dst_prefix = scope_map[i + 1] if scope_map[i + 1].lower() != "none" else "" + + if dst_prefix == "" and (src_prefix + k) in src_state.keys(): + k_src = src_prefix + k + if not k_src.startswith("module."): + logging.info(f"init param, map: {k} from {k_src} in ckpt") + elif ( + k.startswith(dst_prefix) + and k.replace(dst_prefix, src_prefix, 1) in src_state.keys() + ): + k_src = k.replace(dst_prefix, src_prefix, 1) + if not k_src.startswith("module."): + logging.info(f"init param, map: {k} from {k_src} in ckpt") + + if k_src in src_state.keys(): + if ignore_init_mismatch and dst_state[k].shape != src_state[k_src].shape: + logging.info( + f"ignore_init_mismatch:{ignore_init_mismatch}, dst: {k, dst_state[k].shape}, src: {k_src, src_state[k_src].shape}" + ) + else: + dst_state[k] = src_state[k_src] + else: + print(f"Warning, miss key in ckpt: {k}, {path}") + + flag = obj.load_state_dict(dst_state, strict=True) + logging.info(f"Loading ckpt: {path}, status: {flag}") diff --git a/funasr/train_utils/model_summary.py b/funasr/train_utils/model_summary.py new file mode 100644 index 0000000000000000000000000000000000000000..842cd21faee76f7d4c54431bd84162b4261e2ef9 --- /dev/null +++ b/funasr/train_utils/model_summary.py @@ -0,0 +1,72 @@ +import numpy as np +import torch + + +def get_human_readable_count(number: int) -> str: + """Return human_readable_count + + Originated from: + https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/core/memory.py + + Abbreviates an integer number with K, M, B, T for thousands, millions, + billions and trillions, respectively. + Examples: + >>> get_human_readable_count(123) + '123 ' + >>> get_human_readable_count(1234) # (one thousand) + '1 K' + >>> get_human_readable_count(2e6) # (two million) + '2 M' + >>> get_human_readable_count(3e9) # (three billion) + '3 B' + >>> get_human_readable_count(4e12) # (four trillion) + '4 T' + >>> get_human_readable_count(5e15) # (more than trillion) + '5,000 T' + Args: + number: a positive integer number + Return: + A string formatted according to the pattern described above. + """ + assert number >= 0 + labels = [" ", "K", "M", "B", "T"] + num_digits = int(np.floor(np.log10(number)) + 1 if number > 0 else 1) + num_groups = int(np.ceil(num_digits / 3)) + num_groups = min(num_groups, len(labels)) # don't abbreviate beyond trillions + shift = -3 * (num_groups - 1) + number = number * (10**shift) + index = num_groups - 1 + return f"{number:.2f} {labels[index]}" + + +def to_bytes(dtype) -> int: + # torch.float16 -> 16 + return int(str(dtype)[-2:]) // 8 + + +def model_summary(model: torch.nn.Module) -> str: + message = "Model structure:\n" + message += str(model) + + tot_params, num_params = 0, 0 + for name, param in model.named_parameters(): + print( + "name: {}, dtype: {}, device: {}, trainable: {}, shape: {}, numel: {}".format( + name, param.dtype, param.device, param.requires_grad, param.shape, param.numel() + ) + ) + tot_params += param.numel() + if param.requires_grad: + num_params += param.numel() + + percent_trainable = "{:.1f}".format(num_params * 100.0 / tot_params) + tot_params = get_human_readable_count(tot_params) + num_params = get_human_readable_count(num_params) + message += "\n\nModel summary:\n" + message += f" Class Name: {model.__class__.__name__}\n" + message += f" Total Number of model parameters: {tot_params}\n" + message += f" Number of trainable parameters: {num_params} ({percent_trainable}%)\n" + + dtype = next(iter(model.parameters())).dtype + message += f" Type: {dtype}" + return message diff --git a/funasr/train_utils/recursive_op.py b/funasr/train_utils/recursive_op.py new file mode 100644 index 0000000000000000000000000000000000000000..1c2a92fd83b6efd96b9b42d906418c8301c3259c --- /dev/null +++ b/funasr/train_utils/recursive_op.py @@ -0,0 +1,48 @@ +"""Torch utility module.""" + +import torch + +if torch.distributed.is_available(): + from torch.distributed import ReduceOp + + +def recursive_sum(obj, weight: torch.Tensor, distributed: bool = False): + assert weight.dim() == 1, weight.size() + if isinstance(obj, (tuple, list)): + return type(obj)(recursive_sum(v, weight, distributed) for v in obj) + elif isinstance(obj, dict): + return {k: recursive_sum(v, weight, distributed) for k, v in obj.items()} + elif isinstance(obj, torch.Tensor): + assert obj.size() == weight.size(), (obj.size(), weight.size()) + obj = (obj * weight.type(obj.dtype)).sum() + if distributed: + torch.distributed.all_reduce(obj, op=ReduceOp.SUM) + return obj + elif obj is None: + return None + else: + raise ValueError(type(obj)) + + +def recursive_divide(a, b: torch.Tensor): + if isinstance(a, (tuple, list)): + return type(a)(recursive_divide(v, b) for v in a) + elif isinstance(a, dict): + return {k: recursive_divide(v, b) for k, v in a.items()} + elif isinstance(a, torch.Tensor): + assert a.size() == b.size(), (a.size(), b.size()) + return a / b.type(a.dtype) + elif a is None: + return None + else: + raise ValueError(type(a)) + + +def recursive_average(obj, weight: torch.Tensor, distributed: bool = False): + obj = recursive_sum(obj, weight, distributed) + weight = weight.sum() + if distributed: + torch.distributed.all_reduce(weight, op=ReduceOp.SUM) + # Normalize weight to be sum-to-1 + obj = recursive_divide(obj, weight) + return obj, weight diff --git a/funasr/train_utils/set_all_random_seed.py b/funasr/train_utils/set_all_random_seed.py new file mode 100644 index 0000000000000000000000000000000000000000..ebdca3f537aac53bdc6e6cea168c49805bdf2d2f --- /dev/null +++ b/funasr/train_utils/set_all_random_seed.py @@ -0,0 +1,10 @@ +import random + +import numpy as np +import torch + + +def set_all_random_seed(seed: int): + random.seed(seed) + np.random.seed(seed) + torch.random.manual_seed(seed) diff --git a/funasr/train_utils/trainer.py b/funasr/train_utils/trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..d0be9c8070c8174867035c775ff381843a6c77c0 --- /dev/null +++ b/funasr/train_utils/trainer.py @@ -0,0 +1,730 @@ +import math +import os +import time +import torch +import logging +from tqdm import tqdm +from datetime import datetime +import torch.distributed as dist +from torch.cuda.amp import autocast, GradScaler +from contextlib import nullcontext, contextmanager +from pathlib import Path + +from funasr.train_utils.device_funcs import to_device +from funasr.train_utils.recursive_op import recursive_average +from funasr.train_utils.average_nbest_models import average_checkpoints +from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler + +try: + import wandb +except: + wandb = None + + +@contextmanager +def maybe_autocast(enabled): + if enabled: + with autocast(): + yield + else: + yield + + +class Trainer: + """ + A simple trainer class for training a PyTorch model, saving checkpoints at the end of each epoch, + and optionally resuming from a saved checkpoint. + + Attributes: + max_epoch (int): Maximum number of epochs for training. + model (torch.nn.Module): The model to be trained. + optim (torch.optim.Optimizer): The optimizer to use for training. + scheduler (torch.optim.lr_scheduler._LRScheduler): The learning rate scheduler. + dataloader_train (torch.utils.data.DataLoader): DataLoader for the training dataset. + dataloader_val (torch.utils.data.DataLoader): DataLoader for the validation dataset. + output_dir (str): Directory where model checkpoints will be saved. + resume (str, optional): Path to a checkpoint to resume training from. + """ + + def __init__( + self, + local_rank, + use_ddp: bool = False, + use_fsdp: bool = False, + use_fp16: bool = False, + output_dir: str = "./", + **kwargs, + ): + """ + Initializes the Trainer class with the model, optimizer, scheduler, dataloaders, and other settings. + + Args: + model (torch.nn.Module): The model to be trained. + optim (torch.optim.Optimizer): The optimizer to use for training. + scheduler (torch.optim.lr_scheduler._LRScheduler): The learning rate scheduler. + dataloader_train (torch.utils.data.DataLoader): The DataLoader for the training dataset. + dataloader_val (torch.utils.data.DataLoader): The DataLoader for the validation dataset. + **kwargs: Additional keyword arguments: + max_epoch (int): The maximum number of epochs for training. + output_dir (str): The directory where model checkpoints will be saved. Default is './'. + resume (str, optional): The file path to a checkpoint to resume training from. + """ + + self.output_dir = output_dir + if not os.path.exists(self.output_dir): + os.makedirs(self.output_dir, exist_ok=True) + self.resume = kwargs.get("resume", True) + self.start_epoch = 0 + self.max_epoch = kwargs.get("max_epoch", 100) + self.local_rank = local_rank + self.use_ddp = use_ddp + self.use_fsdp = use_fsdp + self.device = kwargs.get("device", "cuda") + # self.kwargs = kwargs + self.log_interval = kwargs.get("log_interval", 50) + self.batch_total = 0 + self.use_fp16 = use_fp16 + self.save_checkpoint_interval = kwargs.get("save_checkpoint_interval", 5000) + self.validate_interval = kwargs.get("validate_interval", -1) + if self.validate_interval < 0: + self.validate_interval = self.save_checkpoint_interval + assert ( + self.save_checkpoint_interval == self.validate_interval + ), f"save_checkpoint_interval must equal to validate_interval" + self.keep_nbest_models = kwargs.get("keep_nbest_models", 500) + self.avg_keep_nbest_models_type = kwargs.get("avg_keep_nbest_models_type", "acc") + self.avg_nbest_model = kwargs.get("avg_nbest_model", 10) + self.accum_grad = kwargs.get("accum_grad", 1) + self.grad_clip = kwargs.get("grad_clip", 10.0) + self.grad_clip_type = kwargs.get("grad_clip_type", 2.0) + + try: + rank = dist.get_rank() + world_size = dist.get_world_size() + except: + rank = 0 + world_size = 1 + logging.warning("distributed is not initialized, only single shard") + self.rank = rank + self.world_size = world_size + self.train_acc_avg = 0.0 + self.train_loss_avg = 0.0 + self.val_acc_avg = 0.0 + self.val_loss_avg = 0.0 + self.best_acc_idx = 0 + self.saved_ckpts = {} + self.step_or_epoch = -1 + self.best_step_or_epoch = "" + self.val_acc_step_or_epoch = {} + self.val_loss_step_or_epoch = {} + + self.reset_gpu_cache = kwargs.get("reset_gpu_cache", False) + self.start_data_split_i = 0 + self.start_step = 0 + self.step_in_epoch = 0 + self.use_wandb = kwargs.get("use_wandb", False) + if self.use_wandb: + wandb.login(key=kwargs.get("wandb_token")) + wandb.init( + config=kwargs, + project=kwargs.get("wandb_project", "my_project"), + entity=kwargs.get("wandb_team", "my_team"), + name=kwargs.get("wandb_exp_name", "my_exp"), + dir=output_dir, + job_type="training", + reinit=True, + ) + + def save_checkpoint( + self, + epoch, + step=None, + model=None, + optim=None, + scheduler=None, + scaler=None, + step_in_epoch=None, + **kwargs, + ): + """ + Saves a checkpoint containing the model's state, the optimizer's state, + and the scheduler's state at the end of the given epoch. This method is + intended to be called at the end of each epoch to save the training progress. + + Args: + epoch (int): The epoch number at which the checkpoint is being saved. + """ + + step_in_epoch = None if step is None else step_in_epoch + if self.rank == 0: + logging.info(f"Save checkpoint: {epoch}, rank: {self.local_rank}\n") + # self.step_or_epoch += 1 + state = { + "epoch": epoch, + "step": step, + "total_step": self.batch_total, + "state_dict": model.state_dict(), + "optimizer": optim.state_dict(), + "scheduler": scheduler.state_dict(), + "saved_ckpts": self.saved_ckpts, + "val_acc_step_or_epoch": self.val_acc_step_or_epoch, + "val_loss_step_or_epoch": self.val_loss_step_or_epoch, + "best_step_or_epoch": self.best_step_or_epoch, + "avg_keep_nbest_models_type": self.avg_keep_nbest_models_type, + "step_in_epoch": step_in_epoch, + "data_split_i": kwargs.get("data_split_i", 0), + "data_split_num": kwargs.get("data_split_num", 1), + "batch_total": self.batch_total, + "train_loss_avg": kwargs.get("train_loss_avg", 0), + "train_acc_avg": kwargs.get("train_acc_avg", 0), + } + step = step_in_epoch + if hasattr(model, "module"): + state["state_dict"] = model.module.state_dict() + + if scaler: + state["scaler_state"] = scaler.state_dict() + + # Create output directory if it does not exist + os.makedirs(self.output_dir, exist_ok=True) + if step is None: + ckpt_name = f"model.pt.ep{epoch}" + else: + ckpt_name = f"model.pt.ep{epoch}.{step}" + filename = os.path.join(self.output_dir, ckpt_name) + torch.save(state, filename) + logging.info(f"Checkpoint saved to {filename}") + + latest = Path(os.path.join(self.output_dir, f"model.pt")) + torch.save(state, latest) + + if self.best_step_or_epoch == "": + self.best_step_or_epoch = ckpt_name + + if self.avg_keep_nbest_models_type == "acc": + if ( + self.val_acc_step_or_epoch[ckpt_name] + >= self.val_acc_step_or_epoch[self.best_step_or_epoch] + ): + self.best_step_or_epoch = ckpt_name + best_ckpt = Path(os.path.join(self.output_dir, f"model.pt.best")) + torch.save(state, best_ckpt) + logging.info( + f"Update best acc: {self.val_acc_step_or_epoch[self.best_step_or_epoch]:.4f}, {best_ckpt}" + ) + else: + logging.info( + f"No improvement in acc: {self.val_acc_step_or_epoch[ckpt_name]:.4f} < {self.val_acc_step_or_epoch[self.best_step_or_epoch]:.4f}, {os.path.join(self.output_dir, self.best_step_or_epoch)}" + ) + elif self.avg_keep_nbest_models_type == "loss": + if ( + self.val_loss_step_or_epoch[ckpt_name] + <= self.val_loss_step_or_epoch[self.best_step_or_epoch] + ): + self.best_step_or_epoch = ckpt_name + best_ckpt = Path(os.path.join(self.output_dir, f"model.pt.best")) + torch.save(state, best_ckpt) + logging.info( + f"Update best loss: {self.val_loss_step_or_epoch[self.best_step_or_epoch]:.4f}, {best_ckpt}" + ) + else: + logging.info( + f"No improvement in loss: {self.val_loss_step_or_epoch[ckpt_name]:.4f} > {self.val_loss_step_or_epoch[self.best_step_or_epoch]:.4f}, {os.path.join(self.output_dir, self.best_step_or_epoch)}" + ) + else: + print("Undo") + self.saved_ckpts[ckpt_name] = getattr( + self, f"val_{self.avg_keep_nbest_models_type}_step_or_epoch" + )[ckpt_name] + if self.keep_nbest_models > 0: + if len(self.saved_ckpts) > self.keep_nbest_models: + if self.avg_keep_nbest_models_type == "acc": + key = min(self.saved_ckpts, key=self.saved_ckpts.get) + else: + key = max(self.saved_ckpts, key=self.saved_ckpts.get) + if key in self.saved_ckpts: + del self.saved_ckpts[key] + filename = os.path.join(self.output_dir, key) + logging.info(f"Delete: {filename}") + if os.path.exists(filename): + os.remove(filename) + + if self.use_ddp or self.use_fsdp: + dist.barrier() + + def resume_checkpoint( + self, + model=None, + optim=None, + scheduler=None, + scaler=None, + ): + """ + Resumes training from a checkpoint at the given file path. + Loads the model's state, the optimizer's state, and the scheduler's state. + + Args: + resume_path (str): The file path to the checkpoint to resume from. + """ + if self.resume: + ckpt = os.path.join(self.output_dir, "model.pt") + if os.path.isfile(ckpt): + checkpoint = torch.load(ckpt, map_location="cpu") + self.start_epoch = checkpoint["epoch"] + # self.model.load_state_dict(checkpoint['state_dict']) + src_state = checkpoint["state_dict"] + dst_state = model.state_dict() + for k in dst_state.keys(): + if not k.startswith("module.") and "module." + k in src_state.keys(): + k_ddp = "module." + k + elif k.startswith("module.") and "module." + k not in src_state.keys(): + k_ddp = k.replace("module.", "", 1) + else: + k_ddp = k + + if k_ddp in src_state.keys(): + dst_state[k] = src_state[k_ddp] + else: + print(f"Miss key in ckpt: model: {k}, ckpt: {k_ddp}") + + model.load_state_dict(dst_state) + optim.load_state_dict(checkpoint["optimizer"]) + scheduler.load_state_dict(checkpoint["scheduler"]) + if scaler is not None and "scaler_state" in checkpoint: + scaler.load_state_dict(checkpoint["scaler_state"]) + + self.saved_ckpts = checkpoint["saved_ckpts"] + self.val_acc_step_or_epoch = ( + checkpoint["val_acc_step_or_epoch"] + if "val_acc_step_or_epoch" in checkpoint + else {} + ) + self.val_loss_step_or_epoch = ( + checkpoint["val_loss_step_or_epoch"] + if "val_loss_step_or_epoch" in checkpoint + else {} + ) + self.best_step_or_epoch = ( + checkpoint["best_step_or_epoch"] if "best_step_or_epoch" in checkpoint else "" + ) + self.start_data_split_i = ( + checkpoint["data_split_i"] if "data_split_i" in checkpoint else 0 + ) + self.batch_total = checkpoint["batch_total"] if "batch_total" in checkpoint else 0 + self.start_step = checkpoint["step"] if "step" in checkpoint else 0 + self.start_step = 0 if self.start_step is None else self.start_step + self.step_in_epoch = ( + checkpoint["step_in_epoch"] if "step_in_epoch" in checkpoint else 0 + ) + self.step_in_epoch = 0 if self.step_in_epoch is None else self.step_in_epoch + print(checkpoint["train_acc_avg"]) + self.train_acc_avg = ( + checkpoint["train_acc_avg"] if "train_acc_avg" in checkpoint else 0 + ) + self.train_loss_avg = ( + checkpoint["train_loss_avg"] if "train_loss_avg" in checkpoint else 0 + ) + model.to(self.device) + print(f"Checkpoint loaded successfully from '{ckpt}'") + else: + print(f"No checkpoint found at '{ckpt}', does not resume status!") + + if self.use_ddp or self.use_fsdp: + dist.barrier() + + def train_epoch( + self, + model=None, + optim=None, + scheduler=None, + scaler=None, + dataloader_train=None, + dataloader_val=None, + epoch=None, + writer=None, + **kwargs, + ): + """ + Defines the training process for a single epoch with gradient accumulation. + Args: + epoch (int): The current epoch number. + """ + if self.use_ddp or self.use_fsdp: + dist.barrier() + logging.info(f"Train epoch: {epoch}, rank: {self.rank}\n") + model.train() + + # Set the number of steps for gradient accumulation + accum_grad = self.accum_grad + # Initialize the gradient accumulation + optim.zero_grad() + speed_stats = {} + + iterator_stop = torch.tensor(0).to(self.device) + + dataloader_train.batch_sampler.set_epoch(epoch) + time_beg = time.perf_counter() + time5 = time_beg + for batch_idx, batch in enumerate(dataloader_train): + # if self.use_ddp or self.use_fsdp: + # dist.all_reduce(iterator_stop, dist.ReduceOp.SUM) + # if iterator_stop > 0: + # break + self.batch_total += 1 + self.step_in_epoch += 1 + time1 = time.perf_counter() + speed_stats["data_load"] = f"{time1-time_beg:0.3f}" + + batch = to_device(batch, self.device) + + my_context = nullcontext + if self.use_ddp or self.use_fsdp: + my_context = model.no_sync if batch_idx % accum_grad != 0 else my_context + with my_context(): + time2 = time.perf_counter() + with maybe_autocast(self.use_fp16): + retval = model(**batch) + + # if ( + # self.reset_gpu_cache + # and (torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024) > 70 + # ): + # torch.cuda.empty_cache() + + loss, stats, weight = retval + stats = {k: v for k, v in stats.items() if v is not None} + if self.use_ddp or self.use_fsdp: + # Apply weighted averaging for loss and stats + loss = (loss * weight.type(loss.dtype)).sum() + # if distributed, this method can also apply all_reduce() + # stats, weight = recursive_average(stats, weight, distributed=True) + if self.use_ddp or self.use_fsdp: + dist.all_reduce(weight, op=dist.ReduceOp.SUM) + # Now weight is summation over all workers + loss /= weight.sum() # shape:[1] -> shape:[] + # Multiply world_size because DistributedDataParallel + # automatically normalizes the gradient by world_size. + loss *= self.world_size + # loss *= self.world_size + # Scale the loss since we're not updating for every mini-batch + loss = loss / accum_grad + + time3 = time.perf_counter() + speed_stats["forward_time"] = f"{time3 - time2:0.3f}" + if self.use_fp16: + scaler.scale(loss).backward() + else: + loss.backward() + time4 = time.perf_counter() + speed_stats["backward_and_AllReaduce_time"] = f"{time4 - time3:0.3f}" + + self.train_loss_avg = ( + self.train_loss_avg * (batch_idx + kwargs.get("start_step", 0)) + + loss.detach().cpu().item() + ) / (batch_idx + kwargs.get("start_step", 0) + 1) + if "acc" in stats: + self.train_acc_avg = ( + self.train_acc_avg * (batch_idx + kwargs.get("start_step", 0)) + + stats["acc"].detach().cpu().item() + ) / (batch_idx + kwargs.get("start_step", 0) + 1) + + # Perform an optimizer step only after accumulating enough gradients + if (batch_idx + 1) % accum_grad == 0: + # Perform gradient clipping if it is set + if self.grad_clip > 0: + grad_norm = torch.nn.utils.clip_grad_norm_( + model.parameters(), + max_norm=self.grad_clip, + norm_type=self.grad_clip_type, + ) + if not torch.isfinite(grad_norm): + logging.warning( + f"The grad norm is {grad_norm}. Skipping updating the model." + ) + optim.zero_grad() # Reset gradients + continue + + # Execute an optimization step (update model parameters) + if self.use_ddp or self.use_fsdp: + dist.barrier() + if self.use_fp16: + scaler.step(optim) + scaler.update() + else: + optim.step() + scheduler.step() + # Clear gradients for the next accumulation stage + optim.zero_grad(set_to_none=True) + + if self.use_ddp or self.use_fsdp: + train_loss_avg = torch.tensor(self.train_loss_avg, dtype=torch.float32).to( + self.device + ) + train_acc_avg = torch.tensor(self.train_acc_avg, dtype=torch.float32).to( + self.device + ) + dist.all_reduce(train_loss_avg, op=dist.ReduceOp.SUM) + dist.all_reduce(train_acc_avg, op=dist.ReduceOp.SUM) + self.train_loss_avg = train_loss_avg.detach().cpu().item() / self.world_size + self.train_acc_avg = train_acc_avg.detach().cpu().item() / self.world_size + + total_time = f"{(time.perf_counter() - time5)/accum_grad:0.3f}" + time5 = time.perf_counter() + + speed_stats["optim_time"] = f"{time5 - time4:0.3f}" + + speed_stats["total_time"] = total_time + lr = scheduler.get_last_lr()[0] + batch_num_epoch = 1 + if hasattr(dataloader_train, "__len__"): + batch_num_epoch = len(dataloader_train) + self.log( + epoch, + batch_idx, + log_step=batch_idx + kwargs.get("start_step", 0), + step_in_epoch=self.step_in_epoch, + batch_num_epoch=batch_num_epoch, + lr=lr, + loss=accum_grad * loss.detach().cpu().item(), + speed_stats=speed_stats, + stats=stats, + writer=writer, + tag="train", + data_split_i=kwargs.get("data_split_i", 0), + data_split_num=kwargs.get("data_split_num", 1), + ) + + if self.step_in_epoch % self.validate_interval == 0: + self.validate_epoch( + model=model, + dataloader_val=dataloader_val, + epoch=epoch, + writer=writer, + step=batch_idx + 1, + step_in_epoch=self.step_in_epoch, + ) + + if self.step_in_epoch % self.save_checkpoint_interval == 0: + self.save_checkpoint( + epoch, + model=model, + optim=optim, + scheduler=scheduler, + scaler=scaler, + step=batch_idx + 1, + step_in_epoch=self.step_in_epoch, + data_split_i=kwargs.get("data_split_i", 0), + data_split_num=kwargs.get("data_split_num", 1), + train_loss_avg=self.train_loss_avg, + train_acc_avg=self.train_acc_avg, + ) + + time_beg = time.perf_counter() + # else: + # if self.use_ddp or self.use_fsdp: + # iterator_stop.fill_(1) + # dist.all_reduce(iterator_stop, dist.ReduceOp.SUM) + + if self.use_ddp or self.use_fsdp: + dist.barrier() + # iterator_stop = torch.tensor(0).to(self.device) + + def validate_epoch( + self, + model=None, + dataloader_val=None, + epoch=None, + writer=None, + **kwargs, + ): + """ + Defines the validation process for a single epoch. + Should be implemented with the actual model validation steps. + + Args: + epoch (int): The current epoch number. + """ + if self.use_ddp or self.use_fsdp: + dist.barrier() + logging.info(f"Validate epoch: {epoch}, rank: {self.rank}\n") + model.eval() + + with torch.no_grad(): + + speed_stats = {} + time5 = time.perf_counter() + iterator_stop = torch.tensor(0).to(self.device) + dataloader_val.batch_sampler.set_epoch(epoch) + for batch_idx, batch in enumerate(dataloader_val): + if self.use_ddp or self.use_fsdp: + dist.all_reduce(iterator_stop, dist.ReduceOp.SUM) + if iterator_stop > 0: + break + time1 = time.perf_counter() + speed_stats["data_load"] = f"{time1 - time5:0.3f}" + batch = to_device(batch, self.device) + + time2 = time.perf_counter() + retval = model(**batch) + time3 = time.perf_counter() + speed_stats["forward_time"] = f"{time3 - time2:0.3f}" + loss, stats, weight = retval + stats = {k: v for k, v in stats.items() if v is not None} + + if self.use_ddp or self.use_fsdp: + # Apply weighted averaging for loss and stats + loss = (loss * weight.type(loss.dtype)).sum() + # if distributed, this method can also apply all_reduce() + # stats, weight = recursive_average(stats, weight, distributed=True) + if self.use_ddp or self.use_fsdp: + dist.all_reduce(weight, op=dist.ReduceOp.SUM) + # Now weight is summation over all workers + loss /= weight.sum() # shape:[1] -> shape:[] + # Multiply world_size because DistributedDataParallel + # automatically normalizes the gradient by world_size. + loss *= self.world_size + + # Scale the loss since we're not updating for every mini-batch + loss = loss + time4 = time.perf_counter() + + if torch.isfinite(loss): + self.val_loss_avg = ( + self.val_loss_avg * batch_idx + loss.detach().cpu().item() + ) / (batch_idx + 1) + + if "acc" in stats: + self.val_acc_avg = ( + self.val_acc_avg * batch_idx + stats["acc"].detach().cpu().item() + ) / (batch_idx + 1) + + if self.use_ddp or self.use_fsdp: + val_loss_avg = torch.tensor(self.val_loss_avg, dtype=torch.float32).to( + self.device + ) + val_acc_avg = torch.tensor(self.val_acc_avg, dtype=torch.float32).to( + self.device + ) + dist.all_reduce(val_loss_avg, op=dist.ReduceOp.SUM) + dist.all_reduce(val_acc_avg, op=dist.ReduceOp.SUM) + self.val_loss_avg = val_loss_avg.detach().cpu().item() / self.world_size + self.val_acc_avg = val_acc_avg.detach().cpu().item() / self.world_size + + time5 = time.perf_counter() + batch_num_epoch = 1 + if hasattr(dataloader_val, "__len__"): + batch_num_epoch = len(dataloader_val) + self.log( + epoch, + batch_idx, + batch_num_epoch=batch_num_epoch, + lr=0.0, + loss=loss.detach().cpu().item(), + speed_stats=speed_stats, + stats=stats, + writer=writer, + tag="val", + ) + + else: + if self.use_ddp or self.use_fsdp: + iterator_stop.fill_(1) + dist.all_reduce(iterator_stop, dist.ReduceOp.SUM) + + if kwargs.get("step_in_epoch", None) is None: + ckpt_name = f"model.pt.ep{epoch}" + else: + ckpt_name = f'model.pt.ep{epoch}.{kwargs.get("step_in_epoch")}' + self.val_acc_step_or_epoch[ckpt_name] = self.val_acc_avg + self.val_loss_step_or_epoch[ckpt_name] = self.val_loss_avg + model.train() + + if self.use_ddp or self.use_fsdp: + dist.barrier() + iterator_stop = torch.tensor(0).to(self.device) + + def log( + self, + epoch=0, + batch_idx=0, + step_in_epoch=0, + batch_num_epoch=-1, + lr=0.0, + loss=0.0, + speed_stats=None, + stats=None, + writer=None, + tag="train", + data_split_i=0, + data_split_num=1, + log_step=None, + **kwargs, + ): + + if (batch_idx + 1) % self.log_interval == 0: + batch_idx = log_step if log_step is not None else batch_idx + gpu_info = ( + "GPU, memory: usage: {:.3f} GB, " + "peak: {:.3f} GB, " + "cache: {:.3f} GB, " + "cache_peak: {:.3f} GB".format( + torch.cuda.memory_allocated() / 1024 / 1024 / 1024, + torch.cuda.max_memory_allocated() / 1024 / 1024 / 1024, + torch.cuda.memory_reserved() / 1024 / 1024 / 1024, + torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, + ) + ) + + loss_avg_epoch = getattr(self, f"{tag}_loss_avg") + acc_avg_epoch = getattr(self, f"{tag}_acc_avg") + description = ( + f"{tag}, " + f"rank: {self.rank}, " + f"epoch: {epoch}/{self.max_epoch}, " + f"data_slice: {data_split_i}/{data_split_num}, " + f"step_in_slice: {batch_idx + 1}/{batch_num_epoch}, step_in_epoch: {step_in_epoch}, total step: {self.batch_total}, " + f"(loss_avg_rank: {loss:.3f}), " + f"(loss_avg_slice: {loss_avg_epoch:.3f}), " + f"(ppl_avg_slice: {math.exp(loss_avg_epoch):.3e}), " + f"(acc_avg_slice: {acc_avg_epoch:.3f}), " + f"(lr: {lr:.3e}), " + f"{[(k, round(v.detach().cpu().item(), 3)) for k, v in stats.items()]}, " + f"{speed_stats}, " + f"{gpu_info}" + ) + logging.info(description) + + description_dict = { + f"rank{self.rank}_loss/{tag}": loss, + f"rank{self.rank}_lr/{tag}": lr, + } + + if writer is not None: + writer.add_scalar(f"rank{self.rank}_loss/{tag}", loss, self.batch_total) + writer.add_scalar(f"rank{self.rank}_lr/{tag}", lr, self.batch_total) + for key, var in stats.items(): + writer.add_scalar( + f"stats_rank{self.rank}_{key}/{tag}", var.item(), self.batch_total + ) + description_dict[f"stats_rank{self.rank}_{key}/{tag}"] = var.item() + for key, var in speed_stats.items(): + writer.add_scalar( + f"stats_rank{self.rank}_{key}/{tag}", eval(var), self.batch_total + ) + description_dict[f"stats_rank{self.rank}_{key}/{tag}"] = eval(var) + if self.use_wandb and wandb is not None: + wandb.log( + description_dict, + setp=self.batch_total, + ) + + def close(self, writer=None): + + if self.use_ddp or self.use_fsdp: + dist.barrier() + + if writer is not None: + writer.close() + + if self.use_ddp or self.use_fsdp: + torch.distributed.destroy_process_group() diff --git a/funasr/train_utils/trainer_ds.py b/funasr/train_utils/trainer_ds.py new file mode 100644 index 0000000000000000000000000000000000000000..0b104da6ee0e148c7a7a56477f4852150a6a9382 --- /dev/null +++ b/funasr/train_utils/trainer_ds.py @@ -0,0 +1,997 @@ +import math +import os +import time +import torch +import logging +from tqdm import tqdm +from datetime import datetime +import torch.distributed as dist +from torch.cuda.amp import autocast, GradScaler +from contextlib import nullcontext, contextmanager +from pathlib import Path +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from funasr.train_utils.device_funcs import to_device +from funasr.train_utils.recursive_op import recursive_average +from funasr.train_utils.average_nbest_models import average_checkpoints +from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler +import funasr.utils.misc as misc_utils + +try: + import wandb +except: + wandb = None + + +@contextmanager +def maybe_autocast(dtype=None, use_deepspeed=False): + if use_deepspeed: + with torch.cuda.amp.autocast(enabled=True, dtype=dtype, cache_enabled=False): + yield + else: + if dtype == torch.float16 or dtype == torch.bfloat16: + yield + # with autocast(enabled=True, dtype=dtype): + # yield + else: + yield + + +class Trainer: + """ + A simple trainer class for training a PyTorch model, saving checkpoints at the end of each epoch, + and optionally resuming from a saved checkpoint. + + Attributes: + max_epoch (int): Maximum number of epochs for training. + model (torch.nn.Module): The model to be trained. + optim (torch.optim.Optimizer): The optimizer to use for training. + scheduler (torch.optim.lr_scheduler._LRScheduler): The learning rate scheduler. + dataloader_train (torch.utils.data.DataLoader): DataLoader for the training dataset. + dataloader_val (torch.utils.data.DataLoader): DataLoader for the validation dataset. + output_dir (str): Directory where model checkpoints will be saved. + resume (str, optional): Path to a checkpoint to resume training from. + """ + + def __init__( + self, + rank=0, + local_rank=0, + world_size=1, + use_ddp: bool = False, + use_fsdp: bool = False, + use_fp16: bool = False, + use_bf16: bool = False, + use_deepspeed: bool = False, + output_dir: str = "./", + **kwargs, + ): + """ + Initializes the Trainer class with the model, optimizer, scheduler, dataloaders, and other settings. + + Args: + model (torch.nn.Module): The model to be trained. + optim (torch.optim.Optimizer): The optimizer to use for training. + scheduler (torch.optim.lr_scheduler._LRScheduler): The learning rate scheduler. + dataloader_train (torch.utils.data.DataLoader): The DataLoader for the training dataset. + dataloader_val (torch.utils.data.DataLoader): The DataLoader for the validation dataset. + **kwargs: Additional keyword arguments: + max_epoch (int): The maximum number of epochs for training. + output_dir (str): The directory where model checkpoints will be saved. Default is './'. + resume (str, optional): The file path to a checkpoint to resume training from. + """ + self.rank = rank + self.local_rank = local_rank + self.world_size = world_size + self.use_ddp = use_ddp + self.use_fsdp = use_fsdp + + self.device = kwargs.get("device", "cuda") + + self.output_dir = output_dir + if not os.path.exists(self.output_dir): + os.makedirs(self.output_dir, exist_ok=True) + self.resume = kwargs.get("resume", True) + self.start_epoch = 0 + self.max_epoch = kwargs.get("max_epoch", 100) + + # self.kwargs = kwargs + self.log_interval = kwargs.get("log_interval", 50) + self.batch_total = 0 + self.dtype = torch.float32 + self.use_fp16 = use_fp16 + self.use_bf16 = use_bf16 + if self.use_fp16: + self.dtype = torch.float16 + if self.use_bf16: + self.dtype = torch.bfloat16 + self.save_checkpoint_interval = kwargs.get("save_checkpoint_interval", 5000) + self.validate_interval = kwargs.get("validate_interval", 5000) + self.keep_nbest_models = kwargs.get("keep_nbest_models", 500) + self.avg_keep_nbest_models_type = kwargs.get("avg_keep_nbest_models_type", "acc") + self.avg_nbest_model = kwargs.get("avg_nbest_model", 10) + self.accum_grad = kwargs.get("accum_grad", 1) + self.grad_clip = kwargs.get("grad_clip", 10.0) + self.grad_clip_type = kwargs.get("grad_clip_type", 2.0) + + self.train_acc_avg = 0.0 + self.train_loss_avg = 0.0 + self.val_acc_avg = 0.0 + self.val_loss_avg = 0.0 + self.best_acc_idx = 0 + self.saved_ckpts = {} + self.step_or_epoch = -1 + self.best_step_or_epoch = "" + self.val_acc_step_or_epoch = {} + self.val_loss_step_or_epoch = {} + + self.reset_gpu_cache = kwargs.get("reset_gpu_cache", False) + self.start_data_split_i = 0 + self.start_step = 0 + self.step_in_epoch = 0 + self.use_wandb = kwargs.get("use_wandb", False) + if self.use_wandb: + wandb.login(key=kwargs.get("wandb_token")) + wandb.init( + config=kwargs, + project=kwargs.get("wandb_project", "my_project"), + entity=kwargs.get("wandb_team", "my_team"), + name=kwargs.get("wandb_exp_name", "my_exp"), + dir=output_dir, + job_type="training", + reinit=True, + ) + tensorboard_dir = os.path.join(output_dir, "tensorboard") + os.makedirs(tensorboard_dir, exist_ok=True) + try: + from tensorboardX import SummaryWriter + + self.writer = SummaryWriter(tensorboard_dir) # if trainer.rank == 0 else None + except: + self.writer = None + + self.use_deepspeed = use_deepspeed + self.deepspeed_config = kwargs.get("deepspeed_config", "") + excludes = kwargs.get("excludes", None) + if excludes is not None: + if isinstance(excludes, str): + excludes = excludes.split(",") + self.excludes = excludes + effective_save_name_excludes = kwargs.get("effective_save_name_excludes", None) + if effective_save_name_excludes is not None: + if isinstance(effective_save_name_excludes, str): + effective_save_name_excludes = effective_save_name_excludes.split(",") + self.effective_save_name_excludes = effective_save_name_excludes + + def save_checkpoint( + self, + epoch, + step=None, + model=None, + optim=None, + scheduler=None, + scaler=None, + step_in_epoch=None, + **kwargs, + ): + """ + Saves a checkpoint containing the model's state, the optimizer's state, + and the scheduler's state at the end of the given epoch. This method is + intended to be called at the end of each epoch to save the training progress. + + Args: + epoch (int): The epoch number at which the checkpoint is being saved. + """ + if self.use_ddp or self.use_fsdp: + dist.barrier() + step_in_epoch = None if step is None else step_in_epoch + if self.use_deepspeed: + + logging.info(f"Save checkpoint: {epoch}, rank: {self.local_rank}\n") + # self.step_or_epoch += 1 + state = { + "epoch": epoch, + # "state_dict": model.state_dict(), + # "optimizer": optim.state_dict(), + # "scheduler": scheduler.state_dict(), + "saved_ckpts": self.saved_ckpts, + "val_acc_step_or_epoch": self.val_acc_step_or_epoch, + "val_loss_step_or_epoch": self.val_loss_step_or_epoch, + "best_step_or_epoch": self.best_step_or_epoch, + "avg_keep_nbest_models_type": self.avg_keep_nbest_models_type, + "step": step, + "step_in_epoch": step_in_epoch, + "data_split_i": kwargs.get("data_split_i", 0), + "data_split_num": kwargs.get("data_split_num", 1), + "batch_total": self.batch_total, + "train_loss_avg": kwargs.get("train_loss_avg", 0), + "train_acc_avg": kwargs.get("train_acc_avg", 0), + } + step = step_in_epoch + if hasattr(model, "module"): + state["state_dict"] = model.module.state_dict() + + if scaler: + state["scaler_state"] = scaler.state_dict() + # Create output directory if it does not exist + os.makedirs(self.output_dir, exist_ok=True) + if step is None: + ckpt_name = f"model.pt.ep{epoch}" + else: + ckpt_name = f"model.pt.ep{epoch}.{step}" + filename = os.path.join(self.output_dir, ckpt_name) + + # torch.save(state, filename) + with torch.no_grad(): + model.save_checkpoint(save_dir=self.output_dir, tag=ckpt_name, client_state=state) + logging.info(f"\nCheckpoint saved to {filename}\n") + latest = Path(os.path.join(self.output_dir, f"model.pt")) + # torch.save(state, latest) + with torch.no_grad(): + model.save_checkpoint(save_dir=self.output_dir, tag=f"model.pt", client_state=state) + if self.best_step_or_epoch == "": + self.best_step_or_epoch = ckpt_name + + if self.avg_keep_nbest_models_type == "acc": + if ( + self.val_acc_step_or_epoch[ckpt_name] + >= self.val_acc_step_or_epoch[self.best_step_or_epoch] + ): + self.best_step_or_epoch = ckpt_name + best_ckpt = Path(os.path.join(self.output_dir, f"model.pt.best")) + # torch.save(state, best_ckpt) + with torch.no_grad(): + model.save_checkpoint( + save_dir=self.output_dir, tag=f"model.pt.best", client_state=state + ) + logging.info( + f"Update best acc: {self.val_acc_step_or_epoch[self.best_step_or_epoch]:.4f}, {best_ckpt}" + ) + else: + logging.info( + f"No improvement in acc: {self.val_acc_step_or_epoch[ckpt_name]:.4f} < {self.val_acc_step_or_epoch[self.best_step_or_epoch]:.4f}, {os.path.join(self.output_dir, self.best_step_or_epoch)}" + ) + elif self.avg_keep_nbest_models_type == "loss": + if ( + self.val_loss_step_or_epoch[ckpt_name] + <= self.val_loss_step_or_epoch[self.best_step_or_epoch] + ): + self.best_step_or_epoch = ckpt_name + best_ckpt = Path(os.path.join(self.output_dir, f"model.pt.best")) + # torch.save(state, best_ckpt) + with torch.no_grad(): + model.save_checkpoint( + save_dir=self.output_dir, tag=f"model.pt.best", client_state=state + ) + logging.info( + f"Update best loss: {self.val_loss_step_or_epoch[self.best_step_or_epoch]:.4f}, {best_ckpt}" + ) + else: + logging.info( + f"No improvement in loss: {self.val_loss_step_or_epoch[ckpt_name]:.4f} > {self.val_loss_step_or_epoch[self.best_step_or_epoch]:.4f}, {os.path.join(self.output_dir, self.best_step_or_epoch)}" + ) + else: + print("Undo") + self.saved_ckpts[ckpt_name] = getattr( + self, f"val_{self.avg_keep_nbest_models_type}_step_or_epoch" + )[ckpt_name] + if self.keep_nbest_models > 0: + if len(self.saved_ckpts) > self.keep_nbest_models: + if self.avg_keep_nbest_models_type == "acc": + key = min(self.saved_ckpts, key=self.saved_ckpts.get) + else: + key = max(self.saved_ckpts, key=self.saved_ckpts.get) + if key in self.saved_ckpts: + del self.saved_ckpts[key] + filename = os.path.join(self.output_dir, key) + logging.info(f"Delete: {filename}") + if os.path.exists(filename): + # os.remove(filename) + misc_utils.smart_remove(filename) + + elif self.use_fsdp: + pass + elif self.rank == 0: + logging.info( + f"Save checkpoint: {epoch}, rank: {self.rank}, local_rank: {self.local_rank}\n" + ) + # self.step_or_epoch += 1 + state = { + "epoch": epoch, + "optimizer": optim.state_dict(), + "scheduler": scheduler.state_dict(), + "saved_ckpts": self.saved_ckpts, + "val_acc_step_or_epoch": self.val_acc_step_or_epoch, + "val_loss_step_or_epoch": self.val_loss_step_or_epoch, + "best_step_or_epoch": self.best_step_or_epoch, + "avg_keep_nbest_models_type": self.avg_keep_nbest_models_type, + "step": step, + "step_in_epoch": step_in_epoch, + "data_split_i": kwargs.get("data_split_i", 0), + "data_split_num": kwargs.get("data_split_num", 1), + "batch_total": self.batch_total, + "train_loss_avg": kwargs.get("train_loss_avg", 0), + "train_acc_avg": kwargs.get("train_acc_avg", 0), + } + step = step_in_epoch + if hasattr(model, "module"): + state_dict = model.module.state_dict() + else: + state_dict = model.state_dict() + + if self.effective_save_name_excludes is not None: + logging.info(f"effective_save_name_excludes: {self.effective_save_name_excludes}") + dst_state_dict = {} + for k in state_dict.keys(): + for k_ex in self.effective_save_name_excludes: + k_tmp = k.replace("module.", "") + if k.startswith(k_ex): + logging.info(f"key: {k} matching: {k_ex}, not save it") + break + else: + dst_state_dict[k] = state_dict[k] + state["state_dict"] = dst_state_dict + else: + state["state_dict"] = state_dict + + if scaler: + state["scaler_state"] = scaler.state_dict() + # Create output directory if it does not exist + os.makedirs(self.output_dir, exist_ok=True) + if step is None: + ckpt_name = f"model.pt.ep{epoch}" + else: + ckpt_name = f"model.pt.ep{epoch}.{step}" + filename = os.path.join(self.output_dir, ckpt_name) + torch.save(state, filename) + + logging.info(f"\nCheckpoint saved to {filename}\n") + latest = Path(os.path.join(self.output_dir, f"model.pt")) + torch.save(state, latest) + if self.best_step_or_epoch == "": + self.best_step_or_epoch = ckpt_name + + if self.avg_keep_nbest_models_type == "acc": + if ( + self.val_acc_step_or_epoch[ckpt_name] + >= self.val_acc_step_or_epoch[self.best_step_or_epoch] + ): + self.best_step_or_epoch = ckpt_name + best_ckpt = Path(os.path.join(self.output_dir, f"model.pt.best")) + torch.save(state, best_ckpt) + logging.info( + f"Update best acc: {self.val_acc_step_or_epoch[self.best_step_or_epoch]:.4f}, {best_ckpt}" + ) + else: + logging.info( + f"No improvement in acc: {self.val_acc_step_or_epoch[ckpt_name]:.4f} < {self.val_acc_step_or_epoch[self.best_step_or_epoch]:.4f}, {os.path.join(self.output_dir, self.best_step_or_epoch)}" + ) + elif self.avg_keep_nbest_models_type == "loss": + if ( + self.val_loss_step_or_epoch[ckpt_name] + <= self.val_loss_step_or_epoch[self.best_step_or_epoch] + ): + self.best_step_or_epoch = ckpt_name + best_ckpt = Path(os.path.join(self.output_dir, f"model.pt.best")) + torch.save(state, best_ckpt) + logging.info( + f"Update best loss: {self.val_loss_step_or_epoch[self.best_step_or_epoch]:.4f}, {best_ckpt}" + ) + else: + logging.info( + f"No improvement in loss: {self.val_loss_step_or_epoch[ckpt_name]:.4f} > {self.val_loss_step_or_epoch[self.best_step_or_epoch]:.4f}, {os.path.join(self.output_dir, self.best_step_or_epoch)}" + ) + else: + print("Undo") + self.saved_ckpts[ckpt_name] = getattr( + self, f"val_{self.avg_keep_nbest_models_type}_step_or_epoch" + )[ckpt_name] + if self.keep_nbest_models > 0: + if len(self.saved_ckpts) > self.keep_nbest_models: + if self.avg_keep_nbest_models_type == "acc": + key = min(self.saved_ckpts, key=self.saved_ckpts.get) + else: + key = max(self.saved_ckpts, key=self.saved_ckpts.get) + if key in self.saved_ckpts: + del self.saved_ckpts[key] + filename = os.path.join(self.output_dir, key) + logging.info(f"Delete: {filename}") + if os.path.exists(filename): + # os.remove(filename) + misc_utils.smart_remove(filename) + + if self.use_ddp or self.use_fsdp: + dist.barrier() + + def resume_checkpoint( + self, + model=None, + optim=None, + scheduler=None, + scaler=None, + ): + """ + Resumes training from a checkpoint at the given file path. + Loads the model's state, the optimizer's state, and the scheduler's state. + + Args: + resume_path (str): The file path to the checkpoint to resume from. + """ + if self.resume: + + if self.use_deepspeed: + ckpt = os.path.join(self.output_dir, "model.pt") + if os.path.exists(ckpt): + _, checkpoint = model.load_checkpoint(self.output_dir, "model.pt") + self.start_epoch = checkpoint["epoch"] + self.saved_ckpts = checkpoint["saved_ckpts"] + self.val_acc_step_or_epoch = ( + checkpoint["val_acc_step_or_epoch"] + if "val_acc_step_or_epoch" in checkpoint + else {} + ) + self.val_loss_step_or_epoch = ( + checkpoint["val_loss_step_or_epoch"] + if "val_loss_step_or_epoch" in checkpoint + else {} + ) + self.best_step_or_epoch = ( + checkpoint["best_step_or_epoch"] + if "best_step_or_epoch" in checkpoint + else "" + ) + self.start_data_split_i = ( + checkpoint["data_split_i"] if "data_split_i" in checkpoint else 0 + ) + self.batch_total = ( + checkpoint["batch_total"] if "batch_total" in checkpoint else 0 + ) + self.start_step = checkpoint["step"] if "step" in checkpoint else 0 + self.start_step = 0 if self.start_step is None else self.start_step + self.step_in_epoch = ( + checkpoint["step_in_epoch"] if "step_in_epoch" in checkpoint else 0 + ) + self.step_in_epoch = 0 if self.step_in_epoch is None else self.step_in_epoch + print(checkpoint["train_acc_avg"]) + self.train_acc_avg = ( + checkpoint["train_acc_avg"] if "train_acc_avg" in checkpoint else 0 + ) + self.train_loss_avg = ( + checkpoint["train_loss_avg"] if "train_loss_avg" in checkpoint else 0 + ) + model.to(self.device) + print(f"Checkpoint loaded successfully from '{ckpt}'") + else: + print(f"No checkpoint found at '{ckpt}', does not resume status!") + else: + + ckpt = os.path.join(self.output_dir, "model.pt") + if os.path.isfile(ckpt): + checkpoint = torch.load(ckpt, map_location="cpu") + self.start_epoch = checkpoint["epoch"] + # self.model.load_state_dict(checkpoint['state_dict']) + src_state = checkpoint["state_dict"] + dst_state = model.state_dict() + for k in dst_state.keys(): + excludes_flag = False + if self.excludes is not None: + for k_ex in self.excludes: + k_tmp = k.replace("module.", "") + if k_tmp.startswith(k_ex): + logging.info(f"key: {k} matching: {k_ex}, excluded") + excludes_flag = True + break + if excludes_flag: + continue + if not k.startswith("module.") and "module." + k in src_state.keys(): + k_ddp = "module." + k + elif k.startswith("module.") and "module." + k not in src_state.keys(): + k_ddp = k.replace("module.", "", 1) + else: + k_ddp = k + if k_ddp in src_state.keys(): + dst_state[k] = src_state[k_ddp] + else: + print(f"Miss key in ckpt: model: {k}, ckpt: {k_ddp}") + + model.load_state_dict(dst_state) + optim.load_state_dict(checkpoint["optimizer"]) + scheduler.load_state_dict(checkpoint["scheduler"]) + if scaler is not None and "scaler_state" in checkpoint: + scaler.load_state_dict(checkpoint["scaler_state"]) + + self.saved_ckpts = checkpoint["saved_ckpts"] + self.val_acc_step_or_epoch = ( + checkpoint["val_acc_step_or_epoch"] + if "val_acc_step_or_epoch" in checkpoint + else {} + ) + self.val_loss_step_or_epoch = ( + checkpoint["val_loss_step_or_epoch"] + if "val_loss_step_or_epoch" in checkpoint + else {} + ) + self.best_step_or_epoch = ( + checkpoint["best_step_or_epoch"] + if "best_step_or_epoch" in checkpoint + else "" + ) + self.start_data_split_i = ( + checkpoint["data_split_i"] if "data_split_i" in checkpoint else 0 + ) + self.batch_total = ( + checkpoint["batch_total"] if "batch_total" in checkpoint else 0 + ) + self.start_step = checkpoint["step"] if "step" in checkpoint else 0 + self.start_step = 0 if self.start_step is None else self.start_step + self.step_in_epoch = ( + checkpoint["step_in_epoch"] if "step_in_epoch" in checkpoint else 0 + ) + self.step_in_epoch = 0 if self.step_in_epoch is None else self.step_in_epoch + print(checkpoint["train_acc_avg"]) + self.train_acc_avg = ( + checkpoint["train_acc_avg"] if "train_acc_avg" in checkpoint else 0 + ) + self.train_loss_avg = ( + checkpoint["train_loss_avg"] if "train_loss_avg" in checkpoint else 0 + ) + model.to(self.device) + print(f"Checkpoint loaded successfully from '{ckpt}'") + else: + print(f"No checkpoint found at '{ckpt}', does not resume status!") + + if self.use_ddp or self.use_fsdp: + dist.barrier() + + def train_epoch( + self, + model=None, + optim=None, + scheduler=None, + scaler=None, + dataloader_train=None, + dataloader_val=None, + epoch=None, + **kwargs, + ): + """ + Defines the training process for a single epoch with gradient accumulation. + Args: + epoch (int): The current epoch number. + """ + if self.use_ddp or self.use_fsdp or self.use_deepspeed: + dist.barrier() + logging.info(f"Train epoch: {epoch}, rank: {self.rank}\n") + model.train() + + # Set the number of steps for gradient accumulation + accum_grad = self.accum_grad + # Initialize the gradient accumulation + optim.zero_grad() + speed_stats = {} + + iterator_stop = torch.tensor(0).to(self.device) + + dataloader_train.batch_sampler.set_epoch(epoch) + time_beg = time.perf_counter() + time5 = time_beg + for batch_idx, batch in enumerate(dataloader_train): + self.batch_total += 1 + self.step_in_epoch += 1 + loss_dict = { + "speed_stats": {}, + "epoch": epoch, + "batch_idx": batch_idx, + "data_split_i": kwargs.get("data_split_i", 0), + "data_split_num": kwargs.get("data_split_num", 1), + "log_step": batch_idx + kwargs.get("start_step", 0), + "batch_total": self.batch_total, + "step_in_epoch": self.step_in_epoch, + } + + time1 = time.perf_counter() + loss_dict["speed_stats"]["data_load"] = f"{time1-time_beg:0.3f}" + + batch = to_device(batch, self.device) + + my_context = nullcontext + if self.use_ddp or self.use_fsdp: + my_context = model.no_sync if batch_idx % accum_grad != 0 else my_context + with my_context(): + time2 = time.perf_counter() + + self.forward_step(model, batch, loss_dict=loss_dict) + + time3 = time.perf_counter() + loss_dict["speed_stats"]["forward_time"] = f"{time3 - time2:0.3f}" + self.backward_step(model, scaler, loss_dict=loss_dict) + + time4 = time.perf_counter() + loss_dict["speed_stats"]["backward_time"] = f"{time4 - time3:0.3f}" + + self.update_step(model, optim, scheduler, scaler, loss_dict=loss_dict) + total_time = f"{(time.perf_counter() - time5):0.3f}" + time5 = time.perf_counter() + + loss_dict["speed_stats"]["optim_time"] = f"{time5 - time4:0.3f}" + + loss_dict["speed_stats"]["total_time"] = total_time + + loss_dict["lr"] = scheduler.get_last_lr()[0] + loss_dict["batch_num_epoch"] = len(dataloader_train) + + self.train_loss_avg = ( + self.train_loss_avg * batch_idx + loss_dict["loss"].detach().cpu().item() + ) / (batch_idx + 1) + if "acc" in loss_dict["stats"]: + self.train_acc_avg = ( + self.train_acc_avg * batch_idx + loss_dict["stats"]["acc"].detach().cpu().item() + ) / (batch_idx + 1) + + self.log(loss_dict, tag="train") + + if self.step_in_epoch % self.validate_interval == 0: + self.validate_epoch( + model=model, + dataloader_val=dataloader_val, + epoch=epoch, + writer=self.writer, + step=batch_idx + 1, + step_in_epoch=self.step_in_epoch, + ) + + if self.step_in_epoch % self.save_checkpoint_interval == 0: + self.save_checkpoint( + epoch, + model=model, + optim=optim, + scheduler=scheduler, + scaler=scaler, + step=batch_idx + 1, + step_in_epoch=self.step_in_epoch, + data_split_i=kwargs.get("data_split_i", 0), + data_split_num=kwargs.get("data_split_num", 1), + train_loss_avg=self.train_loss_avg, + train_acc_avg=self.train_acc_avg, + ) + + time_beg = time.perf_counter() + + if self.use_ddp or self.use_fsdp or self.use_deepspeed: + train_loss_avg = torch.tensor(self.train_loss_avg, dtype=torch.float32).to(self.device) + train_acc_avg = torch.tensor(self.train_acc_avg, dtype=torch.float32).to(self.device) + dist.all_reduce(train_loss_avg, op=dist.ReduceOp.SUM) + dist.all_reduce(train_acc_avg, op=dist.ReduceOp.SUM) + self.train_loss_avg = train_loss_avg.detach().cpu().item() / self.world_size + self.train_acc_avg = train_acc_avg.detach().cpu().item() / self.world_size + + def forward_step(self, model, batch, loss_dict={}): + with maybe_autocast(dtype=self.dtype, use_deepspeed=self.use_deepspeed): + retval = model(**batch) + + loss, stats, weight = retval + stats = {k: v for k, v in stats.items() if v is not None} + + loss_dict["loss"] = loss + loss_dict["stats"] = stats + loss_dict["weight"] = weight + + def backward_step(self, model, scaler, loss_dict={}): + loss = loss_dict["loss"] + + if self.use_deepspeed: + scaled_loss = model.backward(loss) + else: + loss = loss / self.accum_grad + if self.use_fp16 or self.use_bf16: + scaler.scale(loss).backward() + else: + loss.backward() + + def update_step(self, model, optim, scheduler, scaler, loss_dict=None): + batch_idx = loss_dict["batch_idx"] + if self.use_deepspeed: + model.step() + else: + if (batch_idx + 1) % self.accum_grad == 0: + # Perform gradient clipping if it is set + if self.grad_clip > 0: + grad_norm = torch.nn.utils.clip_grad_norm_( + model.parameters(), + max_norm=self.grad_clip, + norm_type=self.grad_clip_type, + ) + if not torch.isfinite(grad_norm): + logging.warning( + f"The grad norm is {grad_norm}. Skipping updating the model." + ) + optim.zero_grad() # Reset gradients + return + + # Execute an optimization step (update model parameters) + if self.use_ddp or self.use_fsdp: + dist.barrier() + if self.use_fp16 or self.use_bf16: + scaler.step(optim) + scaler.update() + else: + optim.step() + scheduler.step() + # Clear gradients for the next accumulation stage + optim.zero_grad(set_to_none=True) + + def validate_epoch( + self, + model=None, + dataloader_val=None, + epoch=None, + writer=None, + **kwargs, + ): + """ + Defines the validation process for a single epoch. + Should be implemented with the actual model validation steps. + + Args: + epoch (int): The current epoch number. + """ + if self.use_ddp or self.use_fsdp or self.use_deepspeed: + dist.barrier() + logging.info(f"Validate epoch: {epoch}, rank: {self.rank}\n") + model.eval() + + with torch.no_grad(): + + speed_stats = {} + time_beg = time.perf_counter() + time5 = time_beg + + dataloader_val.batch_sampler.set_epoch(epoch) + for batch_idx, batch in enumerate(dataloader_val): + + loss_dict = { + "speed_stats": {}, + "epoch": epoch, + "batch_idx": batch_idx, + "data_split_i": kwargs.get("data_split_i", 0), + "data_split_num": kwargs.get("data_split_num", 1), + "log_step": batch_idx + kwargs.get("start_step", 0), + "batch_total": batch_idx + 1, + "step_in_epoch": batch_idx + 1, + "lr": 0.0, + } + + time1 = time.perf_counter() + loss_dict["speed_stats"]["data_load"] = f"{time1 - time_beg:0.3f}" + + batch = to_device(batch, self.device) + + time2 = time.perf_counter() + + self.forward_step(model, batch, loss_dict=loss_dict) + + time3 = time.perf_counter() + loss_dict["speed_stats"]["forward_time"] = f"{time3 - time2:0.3f}" + + total_time = f"{(time.perf_counter() - time5):0.3f}" + time5 = time.perf_counter() + + loss_dict["speed_stats"]["total_time"] = total_time + + loss_dict["batch_num_epoch"] = len(dataloader_val) + + self.log(loss_dict, tag="val") + time_beg = time.perf_counter() + self.val_loss_avg = ( + self.val_loss_avg * batch_idx + loss_dict["loss"].detach().cpu().item() + ) / (batch_idx + 1) + if "acc" in loss_dict["stats"]: + self.val_acc_avg = ( + self.val_acc_avg * batch_idx + + loss_dict["stats"]["acc"].detach().cpu().item() + ) / (batch_idx + 1) + + if self.use_ddp or self.use_fsdp or self.use_deepspeed: + val_loss_avg = torch.tensor(self.val_loss_avg, dtype=torch.float32).to(self.device) + val_acc_avg = torch.tensor(self.val_acc_avg, dtype=torch.float32).to(self.device) + dist.all_reduce(val_loss_avg, op=dist.ReduceOp.SUM) + dist.all_reduce(val_acc_avg, op=dist.ReduceOp.SUM) + self.val_loss_avg = val_loss_avg.detach().cpu().item() / self.world_size + self.val_acc_avg = val_acc_avg.detach().cpu().item() / self.world_size + + if kwargs.get("step_in_epoch", None) is None: + ckpt_name = f"model.pt.ep{epoch}" + else: + ckpt_name = f'model.pt.ep{epoch}.{kwargs.get("step_in_epoch")}' + self.val_acc_step_or_epoch[ckpt_name] = self.val_acc_avg + self.val_loss_step_or_epoch[ckpt_name] = self.val_loss_avg + + if self.use_ddp or self.use_fsdp or self.use_deepspeed: + dist.barrier() + + model.train() + + def log( + self, + loss_dict: dict = None, + tag="train", + **kwargs, + ): + loss = loss_dict["loss"].detach().cpu().item() + epoch = loss_dict["epoch"] + batch_idx = loss_dict["batch_idx"] + step_in_epoch = loss_dict["step_in_epoch"] + batch_total = loss_dict["batch_total"] + batch_num_epoch = loss_dict["batch_num_epoch"] + lr = loss_dict["lr"] + + speed_stats = loss_dict["speed_stats"] + stats = loss_dict["stats"] + data_split_i = loss_dict["data_split_i"] + data_split_num = loss_dict["data_split_num"] + log_step = loss_dict.get("log_step", None) + + if (batch_idx + 1) % self.log_interval == 0: + batch_idx = log_step if log_step is not None else batch_idx + gpu_info = ( + "GPU, memory: usage: {:.3f} GB, " + "peak: {:.3f} GB, " + "cache: {:.3f} GB, " + "cache_peak: {:.3f} GB".format( + torch.cuda.memory_allocated() / 1024 / 1024 / 1024, + torch.cuda.max_memory_allocated() / 1024 / 1024 / 1024, + torch.cuda.memory_reserved() / 1024 / 1024 / 1024, + torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, + ) + ) + + loss_avg_epoch = getattr(self, f"{tag}_loss_avg") + acc_avg_epoch = getattr(self, f"{tag}_acc_avg") + description = ( + f"{tag}, " + f"rank: {self.rank}, " + f"epoch: {epoch}/{self.max_epoch}, " + f"data_slice: {data_split_i}/{data_split_num}, " + f"step_in_slice: {batch_idx + 1}/{batch_num_epoch}, step_in_epoch: {step_in_epoch}, total step: {batch_total}, " + f"(loss_avg_rank: {loss:.3f}), " + f"(loss_avg_slice: {loss_avg_epoch:.3f}), " + f"(ppl_avg_slice: {math.exp(loss_avg_epoch):.3e}), " + f"(acc_avg_slice: {acc_avg_epoch:.3f}), " + f"(lr: {lr:.3e}), " + f"{[(k, round(v.detach().cpu().item(), 3)) for k, v in stats.items()]}, " + f"{speed_stats}, " + f"{gpu_info}" + ) + logging.info(description) + + description_dict = { + f"rank{self.rank}_loss/{tag}": loss, + f"rank{self.rank}_lr/{tag}": lr, + } + + writer = self.writer + if writer is not None: + writer.add_scalar(f"rank{self.rank}_loss/{tag}", loss, batch_total) + writer.add_scalar(f"rank{self.rank}_lr/{tag}", lr, batch_total) + for key, var in stats.items(): + writer.add_scalar(f"stats_rank{self.rank}_{key}/{tag}", var.item(), batch_total) + description_dict[f"stats_rank{self.rank}_{key}/{tag}"] = var.item() + for key, var in speed_stats.items(): + writer.add_scalar(f"stats_rank{self.rank}_{key}/{tag}", eval(var), batch_total) + description_dict[f"stats_rank{self.rank}_{key}/{tag}"] = eval(var) + if self.use_wandb and wandb is not None: + wandb.log( + description_dict, + setp=batch_total, + ) + + def close(self, writer=None): + + if self.use_ddp or self.use_fsdp: + dist.barrier() + + if writer is not None: + writer.close() + + if self.use_ddp or self.use_fsdp: + torch.distributed.destroy_process_group() + + def warp_model(self, model, **kwargs): + + if self.use_deepspeed: + from deepspeed.runtime.zero.stage_1_and_2 import ( + estimate_zero2_model_states_mem_needs_all_live, + ) + from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live + from deepspeed.utils.zero_to_fp32 import convert_zero_checkpoint_to_fp32_state_dict + + local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", 1)) + world_size = int(os.environ.get("WORLD_SIZE", 1)) + + # NOTE(xcsong): look in detail how the memory estimator API works: + # https://deepspeed.readthedocs.io/en/latest/memory.html#discussion + if int(os.environ.get("RANK", 0)) == 0: + logging.info("Estimating model states memory needs (zero2)...") + estimate_zero2_model_states_mem_needs_all_live( + model, + num_gpus_per_node=local_world_size, + num_nodes=world_size // local_world_size, + ) + logging.info("Estimating model states memory needs (zero3)...") + estimate_zero3_model_states_mem_needs_all_live( + model, + num_gpus_per_node=local_world_size, + num_nodes=world_size // local_world_size, + ) + device = None # Init device later + pass # Init DeepSpeed later + + elif self.use_ddp: + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + model = model.cuda(local_rank) + model = DDP( + model, + device_ids=[local_rank], + find_unused_parameters=kwargs.get("train_conf", {}).get( + "find_unused_parameters", False + ), + ) + + else: + model = model.to(device=kwargs.get("device", "cuda")) + + return model + + def warp_optim_scheduler(self, model, **kwargs): + from funasr.optimizers import optim_classes + from funasr.schedulers import scheduler_classes + from omegaconf import OmegaConf, DictConfig + import json + + # optim + logging.info("Build optim") + optim = kwargs.get("optim", "adam") + assert optim in optim_classes + optim_class = optim_classes.get(optim) + optim = optim_class(model.parameters(), **kwargs.get("optim_conf")) + + # scheduler + logging.info("Build scheduler") + scheduler = kwargs.get("scheduler", "warmuplr") + assert scheduler in scheduler_classes + scheduler_class = scheduler_classes.get(scheduler) + scheduler = scheduler_class(optim, **kwargs.get("scheduler_conf")) + + if self.use_deepspeed: + import deepspeed + + args = OmegaConf.create({"deepspeed_config": self.deepspeed_config}) + with open(self.deepspeed_config, "r") as fin: + ds_configs = json.load(fin) + + if "bf16" in ds_configs and ds_configs["bf16"]["enabled"]: + self.dtype = torch.bfloat16 + + if "fp16" in ds_configs and ds_configs["fp16"]["enabled"]: + self.dtype = torch.float16 + if "optimizer" in ds_configs: + # NOTE(xcsong): Disable custom optimizer if it is set in ds_config, + # extremely useful when enable cpu_offload, DeepspeedCpuAdam + # could be 4~5x faster than torch native adam + optim = None + if "scheduler" in ds_configs: + scheduler = None + else: + + def scheduler(opt): + return scheduler_class(opt, **kwargs.get("scheduler_conf")) + + model, optimizer, _, scheduler = deepspeed.initialize( + args=args, + model=model, + optimizer=optim, + lr_scheduler=scheduler, + model_parameters=model.parameters(), + ) + + return model, optim, scheduler diff --git a/funasr/utils/__init__.py b/funasr/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/funasr/utils/compute_det_ctc.py b/funasr/utils/compute_det_ctc.py new file mode 100644 index 0000000000000000000000000000000000000000..47c81821bbcaff7e7ed61a6ab9e3e1e77fd95eb3 --- /dev/null +++ b/funasr/utils/compute_det_ctc.py @@ -0,0 +1,286 @@ +""" This implementation is adapted from https://github.com/wenet-e2e/wekws/blob/main/wekws/bin/compute_det.py.""" + +import os +import json +import logging +import argparse +import threading + +import kaldiio +import torch +from funasr.utils.kws_utils import split_mixed_label + + +class thread_wrapper(threading.Thread): + def __init__(self, func, args=()): + super(thread_wrapper, self).__init__() + self.func = func + self.args = args + self.result = [] + + def run(self): + self.result = self.func(*self.args) + + def get_result(self): + try: + return self.result + except Exception: + return None + + +def space_mixed_label(input_str): + splits = split_mixed_label(input_str) + space_str = ''.join(f'{sub} ' for sub in splits) + return space_str.strip() + + +def read_lists(list_file): + lists = [] + with open(list_file, 'r', encoding='utf8') as fin: + for line in fin: + if line.strip() != '': + lists.append(line.strip()) + return lists + + +def make_pair(wav_lists, trans_lists): + logging.info('make pair for wav-trans list') + + trans_table = {} + for line in trans_lists: + arr = line.strip().replace('\t', ' ').split() + if len(arr) < 2: + logging.debug('invalid line in trans file: {}'.format( + line.strip())) + continue + + trans_table[arr[0]] = line.replace(arr[0],'').strip() + + lists = [] + for line in wav_lists: + arr = line.strip().replace('\t', ' ').split() + if len(arr) == 2 and arr[0] in trans_table: + lists.append( + dict(key=arr[0], + txt=trans_table[arr[0]], + wav=arr[1], + sample_rate=16000)) + else: + logging.debug("can't find corresponding trans for key: {}".format( + arr[0])) + continue + + return lists + + +def count_duration(tid, data_lists): + results = [] + + for obj in data_lists: + assert 'key' in obj + assert 'wav' in obj + assert 'txt' in obj + key = obj['key'] + wav_file = obj['wav'] + txt = obj['txt'] + + try: + rate, waveform = kaldiio.load_mat(wav_file) + waveform = torch.tensor(waveform, dtype=torch.float32) + waveform = waveform.unsqueeze(0) + frames = len(waveform[0]) + duration = frames / float(rate) + except: + logging.info(f'load file failed: {wav_file}') + duration = 0.0 + + obj['duration'] = duration + results.append(obj) + + return results + + +def load_data_and_score(keywords_list, data_file, trans_file, score_file): + # score_table: {uttid: [keywordlist]} + score_table = {} + with open(score_file, 'r', encoding='utf8') as fin: + # read score file and store in table + for line in fin: + arr = line.strip().split() + key = arr[0] + is_detected = arr[1] + if is_detected == 'detected': + if key not in score_table: + score_table.update( + {key: { + 'kw': space_mixed_label(arr[2]), + 'confi': float(arr[3]) + }}) + else: + if key not in score_table: + score_table.update({key: {'kw': 'unknown', 'confi': -1.0}}) + + wav_lists = read_lists(data_file) + trans_lists = read_lists(trans_file) + data_lists = make_pair(wav_lists, trans_lists) + logging.info(f'origin list samples: {len(data_lists)}') + + # count duration for each wave + num_workers = 8 + start = 0 + step = int(len(data_lists) / num_workers) + tasks = [] + for idx in range(num_workers): + if idx != num_workers - 1: + task = thread_wrapper(count_duration, + (idx, data_lists[start:start + step])) + else: + task = thread_wrapper(count_duration, (idx, data_lists[start:])) + task.start() + tasks.append(task) + start += step + + duration_lists = [] + for task in tasks: + task.join() + duration_lists += task.get_result() + logging.info(f'after list samples: {len(duration_lists)}') + + # build empty structure for keyword-filler infos + keyword_filler_table = {} + for keyword in keywords_list: + keyword = space_mixed_label(keyword) + keyword_filler_table[keyword] = {} + keyword_filler_table[keyword]['keyword_table'] = {} + keyword_filler_table[keyword]['keyword_duration'] = 0.0 + keyword_filler_table[keyword]['filler_table'] = {} + keyword_filler_table[keyword]['filler_duration'] = 0.0 + + for obj in duration_lists: + assert 'key' in obj + assert 'wav' in obj + assert 'txt' in obj + assert 'duration' in obj + + key = obj['key'] + wav_file = obj['wav'] + txt = obj['txt'] + txt = space_mixed_label(txt) + txt_regstr_lrblk = ' ' + txt + ' ' + duration = obj['duration'] + assert key in score_table + + for keyword in keywords_list: + keyword = space_mixed_label(keyword) + keyword_regstr_lrblk = ' ' + keyword + ' ' + if txt_regstr_lrblk.find(keyword_regstr_lrblk) != -1: + if keyword == score_table[key]['kw']: + keyword_filler_table[keyword]['keyword_table'].update( + {key: score_table[key]['confi']}) + else: + # uttrance detected but not match this keyword + keyword_filler_table[keyword]['keyword_table'].update( + {key: -1.0}) + keyword_filler_table[keyword]['keyword_duration'] += duration + else: + if keyword == score_table[key]['kw']: + keyword_filler_table[keyword]['filler_table'].update( + {key: score_table[key]['confi']}) + else: + # uttrance if detected, which is not FA for this keyword + keyword_filler_table[keyword]['filler_table'].update( + {key: -1.0}) + keyword_filler_table[keyword]['filler_duration'] += duration + + return keyword_filler_table + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='compute det curve') + parser.add_argument('--keywords', + type=str, + required=True, + help='preset keyword str, input all keywords') + parser.add_argument('--test_data', required=True, help='test data file') + parser.add_argument('--trans_data', + required=True, + default='', + help='transcription of test data') + parser.add_argument('--score_file', required=True, help='score file') + parser.add_argument('--step', + type=float, + default=0.001, + help='threshold step') + parser.add_argument('--stats_dir', + required=True, + help='to save det stats files') + args = parser.parse_args() + + root_logger = logging.getLogger() + handlers = root_logger.handlers[:] + for handler in handlers: + root_logger.removeHandler(handler) + handler.close() + + logging.basicConfig(level=logging.DEBUG, + format='%(asctime)s %(levelname)s %(message)s') + + keywords_list = args.keywords.strip().split(',') + keyword_filler_table = load_data_and_score(keywords_list, args.test_data, + args.trans_data, + args.score_file) + + stats_files = {} + for keyword in keywords_list: + keyword = space_mixed_label(keyword) + keyword_dur = keyword_filler_table[keyword]['keyword_duration'] + keyword_num = len(keyword_filler_table[keyword]['keyword_table']) + filler_dur = keyword_filler_table[keyword]['filler_duration'] + filler_num = len(keyword_filler_table[keyword]['filler_table']) + if keyword_num <= 0: + print('Can\'t compute det for {} without positive sample'.format(keyword)) + continue + if filler_num <= 0: + print('Can\'t compute det for {} without negative sample'.format(keyword)) + continue + + logging.info('Computing det for {}'.format(keyword)) + logging.info(' Keyword duration: {} Hours, wave number: {}'.format( + keyword_dur / 3600.0, keyword_num)) + logging.info(' Filler duration: {} Hours'.format(filler_dur / 3600.0)) + + stats_file = os.path.join(args.stats_dir, 'stats.' + keyword.replace(' ', '_') + '.txt') + with open(stats_file, 'w', encoding='utf8') as fout: + threshold = 0.0 + while threshold <= 1.0: + num_false_reject = 0 + num_true_detect = 0 + # transverse the all keyword_table + for key, confi in keyword_filler_table[keyword][ + 'keyword_table'].items(): + if confi < threshold: + num_false_reject += 1 + else: + num_true_detect += 1 + + num_false_alarm = 0 + # transverse the all filler_table + for key, confi in keyword_filler_table[keyword][ + 'filler_table'].items(): + if confi >= threshold: + num_false_alarm += 1 + # print(f'false alarm: {keyword}, {key}, {confi}') + + # false_reject_rate = num_false_reject / keyword_num + true_detect_rate = num_true_detect / keyword_num + + num_false_alarm = max(num_false_alarm, 1e-6) + false_alarm_per_hour = num_false_alarm / (filler_dur / 3600.0) + false_alarm_rate = num_false_alarm / filler_num + + fout.write('{:.3f} {:.6f} {:.6f} {:.6f}\n'.format( + threshold, true_detect_rate, false_alarm_rate, + false_alarm_per_hour)) + threshold += args.step + + stats_files[keyword] = stats_file diff --git a/funasr/utils/datadir_writer.py b/funasr/utils/datadir_writer.py new file mode 100644 index 0000000000000000000000000000000000000000..882db7708bbe9ea9bb43a0f60c2c2d0d8d560957 --- /dev/null +++ b/funasr/utils/datadir_writer.py @@ -0,0 +1,70 @@ +from pathlib import Path +from typing import Union +import warnings + + +class DatadirWriter: + """Writer class to create kaldi like data directory. + + Examples: + >>> with DatadirWriter("output") as writer: + ... # output/sub.txt is created here + ... subwriter = writer["sub.txt"] + ... # Write "uttidA some/where/a.wav" + ... subwriter["uttidA"] = "some/where/a.wav" + ... subwriter["uttidB"] = "some/where/b.wav" + + """ + + def __init__(self, p: Union[Path, str]): + self.path = Path(p) + self.chilidren = {} + self.fd = None + self.has_children = False + self.keys = set() + + def __enter__(self): + return self + + def __getitem__(self, key: str) -> "DatadirWriter": + if self.fd is not None: + raise RuntimeError("This writer points out a file") + + if key not in self.chilidren: + w = DatadirWriter((self.path / key)) + self.chilidren[key] = w + self.has_children = True + + retval = self.chilidren[key] + return retval + + def __setitem__(self, key: str, value: str): + if self.has_children: + raise RuntimeError("This writer points out a directory") + if key in self.keys: + warnings.warn(f"Duplicated: {key}") + + if self.fd is None: + self.path.parent.mkdir(parents=True, exist_ok=True) + self.fd = self.path.open("w", encoding="utf-8") + + self.keys.add(key) + self.fd.write(f"{key} {value}\n") + self.fd.flush() + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() + + def close(self): + if self.has_children: + prev_child = None + for child in self.chilidren.values(): + child.close() + if prev_child is not None and prev_child.keys != child.keys: + warnings.warn( + f"Ids are mismatching between " f"{prev_child.path} and {child.path}" + ) + prev_child = child + + elif self.fd is not None: + self.fd.close() diff --git a/funasr/utils/dynamic_import.py b/funasr/utils/dynamic_import.py new file mode 100644 index 0000000000000000000000000000000000000000..e0a056bd333faea9ab7f9869e621e16833de8713 --- /dev/null +++ b/funasr/utils/dynamic_import.py @@ -0,0 +1,61 @@ +import importlib.util + +import importlib.util +import inspect +import os.path +import sys + + +def load_module_from_path(file_path): + """ + 从给定的文件路径动态加载模块。 + + :param file_path: 模块文件的绝对路径。 + :return: 加载的模块 + """ + module_name = file_path.split("/")[-1].replace(".py", "") + spec = importlib.util.spec_from_file_location(module_name, file_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def import_module_from_path(file_path: str): + + if file_path.startswith("http"): + from funasr.download.file import download_from_url + + file_path = download_from_url(file_path) + + file_dir = os.path.dirname(file_path) + # file_name = os.path.basename(file_path) + module_name = file_path.split("/")[-1].replace(".py", "") + if len(file_dir) < 1: + file_dir = "./" + sys.path.append(file_dir) + try: + importlib.import_module(module_name) + print(f"Loading remote code successfully: {file_path}") + except Exception as e: + print(f"Loading remote code failed: {file_path}, {e}") + + +# +# def load_module_from_path(module_name, file_path): +# """ +# 从给定的文件路径动态加载模块。 +# +# :param module_name: 动态加载的模块的名称。 +# :param file_path: 模块文件的绝对路径。 +# :return: 加载的模块 +# """ +# # 创建加载模块的spec（规格） +# spec = importlib.util.spec_from_file_location(module_name, file_path) +# +# # 根据spec创建模块 +# module = importlib.util.module_from_spec(spec) +# +# # 执行模块的代码来实际加载它 +# spec.loader.exec_module(module) +# +# return module diff --git a/funasr/utils/export_utils.py b/funasr/utils/export_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ca04d75ce8c3995f079dc2826181bfd3225109ab --- /dev/null +++ b/funasr/utils/export_utils.py @@ -0,0 +1,287 @@ +import os +import torch +import functools + +import warnings + +warnings.filterwarnings("ignore") + + +def export( + model, data_in=None, quantize: bool = False, opset_version: int = 14, type="onnx", **kwargs +): + model_scripts = model.export(**kwargs) + export_dir = kwargs.get("output_dir", os.path.dirname(kwargs.get("init_param"))) + os.makedirs(export_dir, exist_ok=True) + + if not isinstance(model_scripts, (list, tuple)): + model_scripts = (model_scripts,) + for m in model_scripts: + m.eval() + if type == "onnx": + _onnx( + m, + data_in=data_in, + quantize=quantize, + opset_version=opset_version, + export_dir=export_dir, + **kwargs, + ) + elif type == "torchscript": + device = "cuda" if torch.cuda.is_available() else "cpu" + print("Exporting torchscripts on device {}".format(device)) + _torchscripts(m, path=export_dir, device=device) + elif type == "bladedisc": + assert ( + torch.cuda.is_available() + ), "Currently bladedisc optimization for FunASR only supports GPU" + # bladedisc only optimizes encoder/decoder modules + if hasattr(m, "encoder") and hasattr(m, "decoder"): + _bladedisc_opt_for_encdec(m, path=export_dir, enable_fp16=True) + else: + print(f"export_dir: {export_dir}") + _torchscripts(m, path=export_dir, device="cuda") + + elif type == "onnx_fp16": + assert ( + torch.cuda.is_available() + ), "Currently onnx_fp16 optimization for FunASR only supports GPU" + + if hasattr(m, "encoder") and hasattr(m, "decoder"): + _onnx_opt_for_encdec(m, path=export_dir, enable_fp16=True) + + return export_dir + + +def _onnx( + model, + data_in=None, + quantize: bool = False, + opset_version: int = 14, + export_dir: str = None, + **kwargs, +): + + device = kwargs.get("device", "cpu") + dummy_input = model.export_dummy_inputs() + + if isinstance(dummy_input, torch.Tensor): + dummy_input = dummy_input.to(device) + else: + dummy_input = tuple([input.to(device) for input in dummy_input]) + + verbose = kwargs.get("verbose", False) + + if isinstance(model.export_name, str): + export_name = model.export_name + ".onnx" + else: + export_name = model.export_name() + model_path = os.path.join(export_dir, export_name) + torch.onnx.export( + model, + dummy_input, + model_path, + verbose=verbose, + do_constant_folding=True, + opset_version=opset_version, + input_names=model.export_input_names(), + output_names=model.export_output_names(), + dynamic_axes=model.export_dynamic_axes(), + ) + + if quantize: + try: + from onnxruntime.quantization import QuantType, quantize_dynamic + import onnx + except: + raise RuntimeError( + "You are quantizing the onnx model, please install onnxruntime first. via \n`pip install onnx`\n`pip install onnxruntime`." + ) + + quant_model_path = model_path.replace(".onnx", "_quant.onnx") + onnx_model = onnx.load(model_path) + nodes = [n.name for n in onnx_model.graph.node] + nodes_to_exclude = [ + m for m in nodes if "output" in m or "bias_encoder" in m or "bias_decoder" in m + ] + print("Quantizing model from {} to {}".format(model_path, quant_model_path)) + quantize_dynamic( + model_input=model_path, + model_output=quant_model_path, + op_types_to_quantize=["MatMul"], + per_channel=True, + reduce_range=False, + weight_type=QuantType.QUInt8, + nodes_to_exclude=nodes_to_exclude, + ) + + +def _torchscripts(model, path, device="cuda"): + dummy_input = model.export_dummy_inputs() + + if device == "cuda": + model = model.cuda() + if isinstance(dummy_input, torch.Tensor): + dummy_input = dummy_input.cuda() + else: + dummy_input = tuple([i.cuda() for i in dummy_input]) + + model_script = torch.jit.trace(model, dummy_input) + if isinstance(model.export_name, str): + model_script.save(os.path.join(path, f"{model.export_name}".replace("onnx", "torchscript"))) + else: + model_script.save( + os.path.join(path, f"{model.export_name()}".replace("onnx", "torchscript")) + ) + + +def _bladedisc_opt(model, model_inputs, enable_fp16=True): + model = model.eval() + try: + import torch_blade + except Exception as e: + print( + f"Warning, if you are exporting bladedisc, please install it and try it again: pip install -U torch_blade\n" + ) + torch_config = torch_blade.config.Config() + torch_config.enable_fp16 = enable_fp16 + with torch.no_grad(), torch_config: + opt_model = torch_blade.optimize( + model, + allow_tracing=True, + model_inputs=model_inputs, + ) + return opt_model + + +def _rescale_input_hook(m, x, scale): + if len(x) > 1: + return (x[0] / scale, *x[1:]) + else: + return (x[0] / scale,) + + +def _rescale_output_hook(m, x, y, scale): + if isinstance(y, tuple): + return (y[0] / scale, *y[1:]) + else: + return y / scale + + +def _rescale_encoder_model(model, input_data): + # Calculate absmax + absmax = torch.tensor(0).cuda() + + def stat_input_hook(m, x, y): + val = x[0] if isinstance(x, tuple) else x + absmax.copy_(torch.max(absmax, val.detach().abs().max())) + + encoders = model.encoder.model.encoders + hooks = [m.register_forward_hook(stat_input_hook) for m in encoders] + model = model.cuda() + model(*input_data) + for h in hooks: + h.remove() + + # Rescale encoder modules + fp16_scale = int(2 * absmax // 65536) + print(f"rescale encoder modules with factor={fp16_scale}\n\n") + model.encoder.model.encoders0.register_forward_pre_hook( + functools.partial(_rescale_input_hook, scale=fp16_scale), + ) + for name, m in model.encoder.model.named_modules(): + if name.endswith("self_attn"): + m.register_forward_hook(functools.partial(_rescale_output_hook, scale=fp16_scale)) + if name.endswith("feed_forward.w_2"): + state_dict = {k: v / fp16_scale for k, v in m.state_dict().items()} + m.load_state_dict(state_dict) + + +def _bladedisc_opt_for_encdec(model, path, enable_fp16): + # Get input data + # TODO: better to use real data + input_data = model.export_dummy_inputs() + if isinstance(input_data, torch.Tensor): + input_data = input_data.cuda() + else: + input_data = tuple([i.cuda() for i in input_data]) + + # Get input data for decoder module + decoder_inputs = list() + + def get_input_hook(m, x): + decoder_inputs.extend(list(x)) + + hook = model.decoder.register_forward_pre_hook(get_input_hook) + model = model.cuda() + model(*input_data) + hook.remove() + + # Prevent FP16 overflow + if enable_fp16: + _rescale_encoder_model(model, input_data) + + # Export and optimize encoder/decoder modules + model.encoder = _bladedisc_opt(model.encoder, input_data[:2]) + model.decoder = _bladedisc_opt(model.decoder, tuple(decoder_inputs)) + model_script = torch.jit.trace(model, input_data) + model_script.save(os.path.join(path, f"{model.export_name}_blade.torchscript")) + + +def _onnx_opt_for_encdec(model, path, enable_fp16): + + # Get input data + # TODO: better to use real data + input_data = model.export_dummy_inputs() + + if isinstance(input_data, torch.Tensor): + input_data = input_data.cuda() + else: + input_data = tuple([i.cuda() for i in input_data]) + + # Get input data for decoder module + decoder_inputs = list() + + def get_input_hook(m, x): + decoder_inputs.extend(list(x)) + + hook = model.decoder.register_forward_pre_hook(get_input_hook) + model = model.cuda() + model(*input_data) + hook.remove() + + # Prevent FP16 overflow + if enable_fp16: + _rescale_encoder_model(model, input_data) + + fp32_model_path = f"{path}/{model.export_name}_hook.onnx" + print("*" * 50) + print(f"[_onnx_opt_for_encdec(fp32)]: {fp32_model_path}\n\n") + if not os.path.exists(fp32_model_path): + + torch.onnx.export( + model, + input_data, + fp32_model_path, + verbose=False, + do_constant_folding=True, + opset_version=13, + input_names=model.export_input_names(), + output_names=model.export_output_names(), + dynamic_axes=model.export_dynamic_axes(), + ) + + # fp32 to fp16 + fp16_model_path = f"{path}/{model.export_name}_hook_fp16.onnx" + print("*" * 50) + print(f"[_onnx_opt_for_encdec(fp16)]: {fp16_model_path}\n\n") + if os.path.exists(fp32_model_path) and not os.path.exists(fp16_model_path): + try: + from onnxconverter_common import float16 + except: + raise RuntimeError( + "You are converting the onnx model to fp16, please install onnxconverter-common first. via `pip install onnxconverter-common`." + ) + fp32_onnx_model = onnx.load(fp32_model_path) + fp16_onnx_model = float16.convert_float_to_float16(fp32_onnx_model, keep_io_types=True) + onnx.save(fp16_onnx_model, fp16_model_path) diff --git a/funasr/utils/install_model_requirements.py b/funasr/utils/install_model_requirements.py new file mode 100644 index 0000000000000000000000000000000000000000..73a2dcb39cdf6c6698ae8ef61e46d75fb7727fc5 --- /dev/null +++ b/funasr/utils/install_model_requirements.py @@ -0,0 +1,42 @@ +import shutil +import subprocess + + +def install_requirements(requirements_path): + try: + result = pip_install_r(requirements_path) + # check status + if result.returncode == 0: + print("install model requirements successfully") + return True + else: + print("fail to install model requirements! ") + print("error", result.stderr) + return False + except Exception as e: + result = pip_install_r(requirements_path) + # check status + if result.returncode == 0: + print("install model requirements successfully") + return True + else: + print("fail to install model requirements! ") + print("error", result.stderr) + return False + + +def pip_install_r(requirements_path): + cmd = [] + if shutil.which("pip") is not None: + cmd = ["pip"] + elif shutil.which("uv") is not None: + cmd = ["uv", "pip"] + else: + raise RuntimeError("pip not found, failed to install model requirements") + cmd += ["install", "-r", requirements_path] + return subprocess.run( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) diff --git a/funasr/utils/kws_utils.py b/funasr/utils/kws_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..4935040e3f9d0f7b24ba2625a6608789c736b2a4 --- /dev/null +++ b/funasr/utils/kws_utils.py @@ -0,0 +1,284 @@ +import re +import logging + +import torch +import math +from collections import defaultdict +from typing import List, Optional, Tuple + + +symbol_str = '[’!"#$%&\'()*+,-./:;<>=?@，。?★、…【】《》？“”‘’！[\\]^_`{|}~\s]+' + + +def split_mixed_label(input_str): + tokens = [] + s = input_str.lower() + while len(s) > 0: + match = re.match(r'[A-Za-z!?,<>()\']+', s) + if match is not None: + word = match.group(0) + else: + word = s[0:1] + tokens.append(word) + s = s.replace(word, '', 1).strip(' ') + return tokens + + +def query_token_set(txt, symbol_table, lexicon_table): + tokens_str = tuple() + tokens_idx = tuple() + + if txt in symbol_table: + tokens_str = tokens_str + (txt, ) + tokens_idx = tokens_idx + (symbol_table[txt], ) + return tokens_str, tokens_idx + + parts = split_mixed_label(txt) + for part in parts: + if part == '!sil' or part == '(sil)' or part == '': + tokens_str = tokens_str + ('!sil', ) + elif part == '' or part == '': + tokens_str = tokens_str + ('', ) + elif part == '(noise)' or part == 'noise)' or part == '(noise' or part == '': + tokens_str = tokens_str + ('', ) + elif part in symbol_table: + tokens_str = tokens_str + (part, ) + elif part in lexicon_table: + for ch in lexicon_table[part]: + tokens_str = tokens_str + (ch, ) + else: + part = re.sub(symbol_str, '', part) + for ch in part: + tokens_str = tokens_str + (ch, ) + + for ch in tokens_str: + if ch in symbol_table: + tokens_idx = tokens_idx + (symbol_table[ch], ) + elif ch == '!sil': + if 'sil' in symbol_table: + tokens_idx = tokens_idx + (symbol_table['sil'], ) + else: + tokens_idx = tokens_idx + (symbol_table[''], ) + elif ch == '': + if '' in symbol_table: + tokens_idx = tokens_idx + (symbol_table[''], ) + else: + tokens_idx = tokens_idx + (symbol_table[''], ) + else: + if '' in symbol_table: + tokens_idx = tokens_idx + (symbol_table[''], ) + logging.info(f'\'{ch}\' is not in token set, replace with ') + else: + tokens_idx = tokens_idx + (symbol_table[''], ) + logging.info(f'\'{ch}\' is not in token set, replace with ') + + return tokens_str, tokens_idx + + +class KwsCtcPrefixDecoder(): + """Decoder interface wrapper for CTCPrefixDecode.""" + + def __init__( + self, + ctc: torch.nn.Module, + keywords: str, + token_list: list, + seg_dict: dict, + ): + """Initialize class. + + Args: + ctc (torch.nn.Module): The CTC implementation. + For example, :class:`espnet.nets.pytorch_backend.ctc.CTC` + + """ + self.ctc = ctc + self.token_list = token_list + + token_table = {} + for token in token_list: + token_table[token] = token_list.index(token) + + self.keywords_idxset = {0} + self.keywords_token = {} + self.keywords_str = keywords + keywords_list = self.keywords_str.strip().replace(' ', '').split(',') + for keyword in keywords_list: + strs, indexs = query_token_set(keyword, token_table, seg_dict) + self.keywords_token[keyword] = {} + self.keywords_token[keyword]['token_id'] = indexs + self.keywords_token[keyword]['token_str'] = ''.join('%s ' % str(i) for i in indexs) + [ self.keywords_idxset.add(i) for i in indexs ] + + def beam_search( + self, + logits: torch.Tensor, + logits_lengths: torch.Tensor, + keywords_tokenset: set = None, + score_beam_size: int = 3, + path_beam_size: int = 20, + ) -> Tuple[List[List[int]], torch.Tensor]: + """ CTC prefix beam search inner implementation + + Args: + logits (torch.Tensor): (1, max_len, vocab_size) + logits_lengths (torch.Tensor): (1, ) + keywords_tokenset (set): token set for filtering score + score_beam_size (int): beam size for score + path_beam_size (int): beam size for path + + Returns: + List[List[int]]: nbest results + """ + + maxlen = logits.size(0) + ctc_probs = logits + cur_hyps = [(tuple(), (1.0, 0.0, []))] + + # CTC beam search step by step + for t in range(0, maxlen): + probs = ctc_probs[t] # (vocab_size,) + # key: prefix, value (pb, pnb), default value(-inf, -inf) + next_hyps = defaultdict(lambda: (0.0, 0.0, [])) + + # 2.1 First beam prune: select topk best + top_k_probs, top_k_index = probs.topk( + score_beam_size) # (score_beam_size,) + + # filter prob score that is too small + filter_probs = [] + filter_index = [] + for prob, idx in zip(top_k_probs.tolist(), top_k_index.tolist()): + if keywords_tokenset is not None: + if prob > 0.05 and idx in keywords_tokenset: + filter_probs.append(prob) + filter_index.append(idx) + else: + if prob > 0.05: + filter_probs.append(prob) + filter_index.append(idx) + + if len(filter_index) == 0: + continue + + for s in filter_index: + ps = probs[s].item() + # print(f'frame:{t}, token:{s}, score:{ps}') + + for prefix, (pb, pnb, cur_nodes) in cur_hyps: + last = prefix[-1] if len(prefix) > 0 else None + if s == 0: # blank + n_pb, n_pnb, nodes = next_hyps[prefix] + n_pb = n_pb + pb * ps + pnb * ps + nodes = cur_nodes.copy() + next_hyps[prefix] = (n_pb, n_pnb, nodes) + elif s == last: + if not math.isclose(pnb, 0.0, abs_tol=0.000001): + # Update *ss -> *s; + n_pb, n_pnb, nodes = next_hyps[prefix] + n_pnb = n_pnb + pnb * ps + nodes = cur_nodes.copy() + if ps > nodes[-1]['prob']: # update frame and prob + nodes[-1]['prob'] = ps + nodes[-1]['frame'] = t + next_hyps[prefix] = (n_pb, n_pnb, nodes) + + if not math.isclose(pb, 0.0, abs_tol=0.000001): + # Update *s-s -> *ss, - is for blank + n_prefix = prefix + (s, ) + n_pb, n_pnb, nodes = next_hyps[n_prefix] + n_pnb = n_pnb + pb * ps + nodes = cur_nodes.copy() + nodes.append(dict(token=s, frame=t, + prob=ps)) # to record token prob + next_hyps[n_prefix] = (n_pb, n_pnb, nodes) + else: + n_prefix = prefix + (s, ) + n_pb, n_pnb, nodes = next_hyps[n_prefix] + if nodes: + if ps > nodes[-1]['prob']: # update frame and prob + nodes[-1]['prob'] = ps + nodes[-1]['frame'] = t + else: + nodes = cur_nodes.copy() + nodes.append(dict(token=s, frame=t, + prob=ps)) # to record token prob + n_pnb = n_pnb + pb * ps + pnb * ps + next_hyps[n_prefix] = (n_pb, n_pnb, nodes) + + # 2.2 Second beam prune + next_hyps = sorted(next_hyps.items(), + key=lambda x: (x[1][0] + x[1][1]), + reverse=True) + + cur_hyps = next_hyps[:path_beam_size] + + hyps = [(y[0], y[1][0] + y[1][1], y[1][2]) for y in cur_hyps] + return hyps + + + def is_sublist(self, main_list, check_list): + if len(main_list) < len(check_list): + return -1 + + if len(main_list) == len(check_list): + return 0 if main_list == check_list else -1 + + for i in range(len(main_list) - len(check_list)): + if main_list[i] == check_list[0]: + for j in range(len(check_list)): + if main_list[i + j] != check_list[j]: + break + else: + return i + else: + return -1 + + + def _decode_inside( + self, + logits: torch.Tensor, + logits_lengths: torch.Tensor, + ): + hyps = self.beam_search(logits, logits_lengths, self.keywords_idxset) + + hit_keyword = None + hit_score = 1.0 + # start = 0; end = 0 + for one_hyp in hyps: + prefix_ids = one_hyp[0] + # path_score = one_hyp[1] + prefix_nodes = one_hyp[2] + assert len(prefix_ids) == len(prefix_nodes) + for word in self.keywords_token.keys(): + lab = self.keywords_token[word]['token_id'] + offset = self.is_sublist(prefix_ids, lab) + if offset != -1: + hit_keyword = word + for idx in range(offset, offset + len(lab)): + hit_score *= prefix_nodes[idx]['prob'] + break + if hit_keyword is not None: + hit_score = math.sqrt(hit_score) + break + + if hit_keyword is not None: + return True, hit_keyword, hit_score + else: + return False, None, None + + + def decode(self, x: torch.Tensor): + """Get an initial state for decoding. + + Args: + x (torch.Tensor): The encoded feature tensor + + Returns: decode result + + """ + + raw_logp = self.ctc.softmax(x.unsqueeze(0)).detach().squeeze(0).cpu() + xlen = torch.tensor([raw_logp.size(1)]) + + return self._decode_inside(raw_logp, xlen) diff --git a/funasr/utils/load_utils.py b/funasr/utils/load_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ccd8545e146ed5efa5c0654ab561b4acce938c40 --- /dev/null +++ b/funasr/utils/load_utils.py @@ -0,0 +1,262 @@ +import os +import torch +import json +from io import BytesIO +import torch.distributed as dist +import numpy as np +import kaldiio +import librosa +import torchaudio +import time +import logging +from torch.nn.utils.rnn import pad_sequence + +try: + from funasr.download.file import download_from_url +except: + print("urllib is not installed, if you infer from url, please install it first.") +import pdb +import subprocess +from subprocess import CalledProcessError, run + +try: + from pydub import AudioSegment +except: + pass + + +def is_ffmpeg_installed(): + try: + output = subprocess.check_output(["ffmpeg", "-version"], stderr=subprocess.STDOUT) + return "ffmpeg version" in output.decode("utf-8") + except (subprocess.CalledProcessError, FileNotFoundError): + return False + + +use_ffmpeg = False +if is_ffmpeg_installed(): + use_ffmpeg = True +else: + print( + "Notice: ffmpeg is not installed. torchaudio is used to load audio\n" + "If you want to use ffmpeg backend to load audio, please install it by:" + "\n\tsudo apt install ffmpeg # ubuntu" + "\n\t# brew install ffmpeg # mac" + ) + + +def load_audio_text_image_video( + data_or_path_or_list, + fs: int = 16000, + audio_fs: int = 16000, + data_type="sound", + tokenizer=None, + **kwargs, +): + if isinstance(data_or_path_or_list, (list, tuple)): + if data_type is not None and isinstance(data_type, (list, tuple)): + data_types = [data_type] * len(data_or_path_or_list) + data_or_path_or_list_ret = [[] for d in data_type] + for i, (data_type_i, data_or_path_or_list_i) in enumerate( + zip(data_types, data_or_path_or_list) + ): + for j, (data_type_j, data_or_path_or_list_j) in enumerate( + zip(data_type_i, data_or_path_or_list_i) + ): + data_or_path_or_list_j = load_audio_text_image_video( + data_or_path_or_list_j, + fs=fs, + audio_fs=audio_fs, + data_type=data_type_j, + tokenizer=tokenizer, + **kwargs, + ) + data_or_path_or_list_ret[j].append(data_or_path_or_list_j) + + return data_or_path_or_list_ret + else: + return [ + load_audio_text_image_video( + audio, fs=fs, audio_fs=audio_fs, data_type=data_type, **kwargs + ) + for audio in data_or_path_or_list + ] + if isinstance(data_or_path_or_list, str) and data_or_path_or_list.startswith( + ("http://", "https://") + ): # download url to local file + data_or_path_or_list = download_from_url(data_or_path_or_list) + + if (isinstance(data_or_path_or_list, str) and os.path.exists(data_or_path_or_list)) or hasattr(data_or_path_or_list, 'read'): # local file or bytes io + if data_type is None or data_type == "sound": + if hasattr(data_or_path_or_list, "read") and hasattr(data_or_path_or_list, "seek"): + data_or_path_or_list.seek(0) + # if use_ffmpeg: + # data_or_path_or_list = _load_audio_ffmpeg(data_or_path_or_list, sr=fs) + # data_or_path_or_list = torch.from_numpy(data_or_path_or_list).squeeze() # [n_samples,] + # else: + # data_or_path_or_list, audio_fs = torchaudio.load(data_or_path_or_list) + # if kwargs.get("reduce_channels", True): + # data_or_path_or_list = data_or_path_or_list.mean(0) + try: + data_or_path_or_list, audio_fs = torchaudio.load(data_or_path_or_list) + if kwargs.get("reduce_channels", True): + data_or_path_or_list = data_or_path_or_list.mean(0) + except: + data_or_path_or_list = _load_audio_ffmpeg(data_or_path_or_list, sr=fs) + data_or_path_or_list = torch.from_numpy( + data_or_path_or_list + ).squeeze() # [n_samples,] + elif data_type == "text" and tokenizer is not None: + data_or_path_or_list = tokenizer.encode(data_or_path_or_list) + elif data_type == "image": # undo + pass + elif data_type == "video": # undo + pass + + # if data_in is a file or url, set is_final=True + if "cache" in kwargs: + kwargs["cache"]["is_final"] = True + kwargs["cache"]["is_streaming_input"] = False + elif isinstance(data_or_path_or_list, str) and data_type == "text" and tokenizer is not None: + data_or_path_or_list = tokenizer.encode(data_or_path_or_list) + elif isinstance(data_or_path_or_list, np.ndarray): # audio sample point + data_or_path_or_list = torch.from_numpy(data_or_path_or_list) # .squeeze() # [n_samples,] + elif isinstance(data_or_path_or_list, str) and data_type == "kaldi_ark": + data_mat = kaldiio.load_mat(data_or_path_or_list) + if isinstance(data_mat, tuple): + audio_fs, mat = data_mat + else: + mat = data_mat + if mat.dtype == "int16" or mat.dtype == "int32": + mat = mat.astype(np.float64) + mat = mat / 32768 + if mat.ndim == 2: + mat = mat[:, 0] + data_or_path_or_list = mat + else: + pass + # print(f"unsupport data type: {data_or_path_or_list}, return raw data") + + if audio_fs != fs and data_type != "text": + resampler = torchaudio.transforms.Resample(audio_fs, fs) + data_or_path_or_list = resampler(data_or_path_or_list[None, :])[0, :] + return data_or_path_or_list + + +def load_bytes(input): + """ + try: + input = validate_frame_rate(input) + except: + pass + """ + middle_data = np.frombuffer(input, dtype=np.int16) + middle_data = np.asarray(middle_data) + if middle_data.dtype.kind not in "iu": + raise TypeError("'middle_data' must be an array of integers") + dtype = np.dtype("float32") + if dtype.kind != "f": + raise TypeError("'dtype' must be a floating point type") + + i = np.iinfo(middle_data.dtype) + abs_max = 2 ** (i.bits - 1) + offset = i.min + abs_max + array = np.frombuffer((middle_data.astype(dtype) - offset) / abs_max, dtype=np.float32) + return array + + +def validate_frame_rate( + input, + fs: int = 16000, +): + + # 将文件读取为字节流 + byte_data = BytesIO(input) + + # 使用 pydub 加载音频 + try: + audio = AudioSegment.from_file(byte_data) + except: + raise RuntimeError( + "You are decoding the pcm data, please install pydub first. via `pip install pydub`." + ) + + # 确保采样率为 16000 Hz + if audio.frame_rate != fs: + audio = audio.set_frame_rate(fs) + + # 将重新采样后的音频导出为字节流 + output = BytesIO() + audio.export(output, format="wav") + output.seek(0) + + # 获取重新采样后的字节流数据 + input = output.read() + + return input + + +def extract_fbank(data, data_len=None, data_type: str = "sound", frontend=None, **kwargs): + if isinstance(data, np.ndarray): + data = torch.from_numpy(data) + if len(data.shape) < 2: + data = data[None, :] # data: [batch, N] + data_len = [data.shape[1]] if data_len is None else data_len + elif isinstance(data, torch.Tensor): + if len(data.shape) < 2: + data = data[None, :] # data: [batch, N] + data_len = [data.shape[1]] if data_len is None else data_len + elif isinstance(data, (list, tuple)): + data_list, data_len = [], [] + for data_i in data: + if isinstance(data_i, np.ndarray): + data_i = torch.from_numpy(data_i) + data_list.append(data_i) + data_len.append(data_i.shape[0]) + data = pad_sequence(data_list, batch_first=True) # data: [batch, N] + + data, data_len = frontend(data, data_len, **kwargs) + + if isinstance(data_len, (list, tuple)): + data_len = torch.tensor([data_len]) + return data.to(torch.float32), data_len.to(torch.int32) + + +def _load_audio_ffmpeg(file: str, sr: int = 16000): + """ + Open an audio file and read as mono waveform, resampling as necessary + + Parameters + ---------- + file: str + The audio file to open + + sr: int + The sample rate to resample the audio if necessary + + Returns + ------- + A NumPy array containing the audio waveform, in float32 dtype. + """ + + # This launches a subprocess to decode audio while down-mixing + # and resampling as necessary. Requires the ffmpeg CLI in PATH. + # fmt: off + cmd = [ + "ffmpeg", + "-nostdin", + "-threads", "0", + "-i", file, + "-f", "s16le", + "-ac", "1", + "-acodec", "pcm_s16le", + "-ar", str(sr), + "-" + ] + # fmt: on + try: + out = run(cmd, capture_output=True, check=True).stdout + except CalledProcessError as e: + raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e + + return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0 diff --git a/funasr/utils/misc.py b/funasr/utils/misc.py new file mode 100644 index 0000000000000000000000000000000000000000..eb17f9723c9084c8a00a4f0e2e403aba2ace10ea --- /dev/null +++ b/funasr/utils/misc.py @@ -0,0 +1,119 @@ +import os +import io +import shutil +import logging +from collections import OrderedDict +import numpy as np +from omegaconf import DictConfig, OmegaConf + + +def statistic_model_parameters(model, prefix=None): + var_dict = model.state_dict() + numel = 0 + for i, key in enumerate( + sorted(list([x for x in var_dict.keys() if "num_batches_tracked" not in x])) + ): + if prefix is None or key.startswith(prefix): + numel += var_dict[key].numel() + return numel + + +def int2vec(x, vec_dim=8, dtype=np.int32): + b = ("{:0" + str(vec_dim) + "b}").format(x) + # little-endian order: lower bit first + return (np.array(list(b)[::-1]) == "1").astype(dtype) + + +def seq2arr(seq, vec_dim=8): + return np.row_stack([int2vec(int(x), vec_dim) for x in seq]) + + +def load_scp_as_dict(scp_path, value_type="str", kv_sep=" "): + with io.open(scp_path, "r", encoding="utf-8") as f: + ret_dict = OrderedDict() + for one_line in f.readlines(): + one_line = one_line.strip() + pos = one_line.find(kv_sep) + key, value = one_line[:pos], one_line[pos + 1 :] + if value_type == "list": + value = value.split(" ") + ret_dict[key] = value + return ret_dict + + +def load_scp_as_list(scp_path, value_type="str", kv_sep=" "): + with io.open(scp_path, "r", encoding="utf8") as f: + ret_dict = [] + for one_line in f.readlines(): + one_line = one_line.strip() + pos = one_line.find(kv_sep) + key, value = one_line[:pos], one_line[pos + 1 :] + if value_type == "list": + value = value.split(" ") + ret_dict.append((key, value)) + return ret_dict + + +def deep_update(original, update): + for key, value in update.items(): + if isinstance(value, dict) and key in original: + if len(value) == 0: + original[key] = value + deep_update(original[key], value) + else: + original[key] = value + + +def prepare_model_dir(**kwargs): + + os.makedirs(kwargs.get("output_dir", "./"), exist_ok=True) + + yaml_file = os.path.join(kwargs.get("output_dir", "./"), "config.yaml") + OmegaConf.save(config=kwargs, f=yaml_file) + logging.info(f"kwargs: {kwargs}") + logging.info("config.yaml is saved to: %s", yaml_file) + + model_path = kwargs.get("model_path", None) + if model_path is not None: + config_json = os.path.join(model_path, "configuration.json") + if os.path.exists(config_json): + shutil.copy( + config_json, os.path.join(kwargs.get("output_dir", "./"), "configuration.json") + ) + + +def extract_filename_without_extension(file_path): + """ + 从给定的文件路径中提取文件名（不包含路径和扩展名） + :param file_path: 完整的文件路径 + :return: 文件名（不含路径和扩展名） + """ + # 首先，使用os.path.basename获取路径中的文件名部分（含扩展名） + filename_with_extension = os.path.basename(file_path) + # 然后，使用os.path.splitext分离文件名和扩展名 + filename, extension = os.path.splitext(filename_with_extension) + # 返回不包含扩展名的文件名 + return filename + + +def smart_remove(path): + """Intelligently removes files, empty directories, and non-empty directories recursively.""" + # Check if the provided path exists + if not os.path.exists(path): + print(f"{path} does not exist.") + return + + # If the path is a file, delete it + if os.path.isfile(path): + os.remove(path) + print(f"File {path} has been deleted.") + # If the path is a directory + elif os.path.isdir(path): + try: + # Attempt to remove an empty directory + os.rmdir(path) + print(f"Empty directory {path} has been deleted.") + except OSError: + # If the directory is not empty, remove it along with all its contents + shutil.rmtree(path) + print(f"Non-empty directory {path} has been recursively deleted.") diff --git a/funasr/utils/postprocess_utils.py b/funasr/utils/postprocess_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..60ebb56fd77c38edf5fe481e246ec3cb65cd16de --- /dev/null +++ b/funasr/utils/postprocess_utils.py @@ -0,0 +1,423 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +import string +import logging +from typing import Any, List, Union + + +def isChinese(ch: str): + if "\u4e00" <= ch <= "\u9fff" or "\u0030" <= ch <= "\u0039" or ch == "@": + return True + return False + + +def isAllChinese(word: Union[List[Any], str]): + word_lists = [] + for i in word: + cur = i.replace(" ", "") + cur = cur.replace("", "") + cur = cur.replace("", "") + cur = cur.replace("", "") + cur = cur.replace("", "") + word_lists.append(cur) + + if len(word_lists) == 0: + return False + + for ch in word_lists: + if isChinese(ch) is False: + return False + return True + + +def isAllAlpha(word: Union[List[Any], str]): + word_lists = [] + for i in word: + cur = i.replace(" ", "") + cur = cur.replace("", "") + cur = cur.replace("", "") + cur = cur.replace("", "") + cur = cur.replace("", "") + word_lists.append(cur) + + if len(word_lists) == 0: + return False + + for ch in word_lists: + if ch.isalpha() is False and ch != "'": + return False + elif ch.isalpha() is True and isChinese(ch) is True: + return False + + return True + + +# def abbr_dispose(words: List[Any]) -> List[Any]: +def abbr_dispose(words: List[Any], time_stamp: List[List] = None) -> List[Any]: + words_size = len(words) + word_lists = [] + abbr_begin = [] + abbr_end = [] + last_num = -1 + ts_lists = [] + ts_nums = [] + ts_index = 0 + for num in range(words_size): + if num <= last_num: + continue + + if len(words[num]) == 1 and words[num].encode("utf-8").isalpha(): + if ( + num + 1 < words_size + and words[num + 1] == " " + and num + 2 < words_size + and len(words[num + 2]) == 1 + and words[num + 2].encode("utf-8").isalpha() + ): + # found the begin of abbr + abbr_begin.append(num) + num += 2 + abbr_end.append(num) + # to find the end of abbr + while True: + num += 1 + if num < words_size and words[num] == " ": + num += 1 + if ( + num < words_size + and len(words[num]) == 1 + and words[num].encode("utf-8").isalpha() + ): + abbr_end.pop() + abbr_end.append(num) + last_num = num + else: + break + else: + break + + for num in range(words_size): + if words[num] == " ": + ts_nums.append(ts_index) + else: + ts_nums.append(ts_index) + ts_index += 1 + last_num = -1 + for num in range(words_size): + if num <= last_num: + continue + + if num in abbr_begin: + if time_stamp is not None: + begin = time_stamp[ts_nums[num]][0] + abbr_word = words[num].upper() + num += 1 + while num < words_size: + if num in abbr_end: + abbr_word += words[num].upper() + last_num = num + break + else: + if words[num].encode("utf-8").isalpha(): + abbr_word += words[num].upper() + num += 1 + word_lists.append(abbr_word) + if time_stamp is not None and ts_nums[num] < len(time_stamp): + end = time_stamp[ts_nums[num]][1] + ts_lists.append([begin, end]) + else: + word_lists.append(words[num]) + # length of time_stamp may not equal to length of words because of the (somehow improper) threshold set in timestamp_tools.py line 46, e.g., length of time_stamp can be zero but length of words is not. + # Moreover, move "word_lists.append(words[num])" into if clause, to keep length of word_lists and length of ts_lists equal. + if time_stamp is not None and ts_nums[num] < len(time_stamp) and words[num] != " ": + begin = time_stamp[ts_nums[num]][0] + end = time_stamp[ts_nums[num]][1] + ts_lists.append([begin, end]) + begin = end + + if time_stamp is not None: + return word_lists, ts_lists + else: + return word_lists + + +def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): + middle_lists = [] + word_lists = [] + word_item = "" + ts_lists = [] + + # wash words lists + for i in words: + word = "" + if isinstance(i, str): + word = i + else: + word = i.decode("utf-8") + + if word in ["", "", "", ""]: + continue + else: + middle_lists.append(word) + + # all chinese characters + if isAllChinese(middle_lists): + for i, ch in enumerate(middle_lists): + word_lists.append(ch.replace(" ", "")) + if time_stamp is not None: + ts_lists = time_stamp + + # all alpha characters + elif isAllAlpha(middle_lists): + ts_flag = True + for i, ch in enumerate(middle_lists): + if ts_flag and time_stamp is not None: + begin = time_stamp[i][0] + end = time_stamp[i][1] + word = "" + if "@@" in ch: + word = ch.replace("@@", "") + word_item += word + if time_stamp is not None: + ts_flag = False + end = time_stamp[i][1] + else: + word_item += ch + word_lists.append(word_item) + word_lists.append(" ") + word_item = "" + if time_stamp is not None: + ts_flag = True + end = time_stamp[i][1] + ts_lists.append([begin, end]) + begin = end + + # mix characters + else: + alpha_blank = False + ts_flag = True + begin = -1 + end = -1 + for i, ch in enumerate(middle_lists): + if ts_flag and time_stamp is not None: + begin = time_stamp[i][0] + end = time_stamp[i][1] + word = "" + if isAllChinese(ch): + if alpha_blank is True: + word_lists.pop() + word_lists.append(ch) + alpha_blank = False + if time_stamp is not None: + ts_flag = True + ts_lists.append([begin, end]) + begin = end + elif "@@" in ch: + word = ch.replace("@@", "") + word_item += word + alpha_blank = False + if time_stamp is not None: + ts_flag = False + end = time_stamp[i][1] + elif isAllAlpha(ch): + word_item += ch + word_lists.append(word_item) + word_lists.append(" ") + word_item = "" + alpha_blank = True + if time_stamp is not None: + ts_flag = True + end = time_stamp[i][1] + ts_lists.append([begin, end]) + begin = end + else: + word_lists.append(ch) + + if time_stamp is not None: + word_lists, ts_lists = abbr_dispose(word_lists, ts_lists) + real_word_lists = [] + for ch in word_lists: + if ch != " ": + real_word_lists.append(ch) + sentence = " ".join(real_word_lists).strip() + return sentence, ts_lists, real_word_lists + else: + word_lists = abbr_dispose(word_lists) + real_word_lists = [] + for ch in word_lists: + if ch != " ": + real_word_lists.append(ch) + sentence = "".join(word_lists).strip() + return sentence, real_word_lists + + +def sentence_postprocess_sentencepiece(words): + middle_lists = [] + word_lists = [] + word_item = "" + + # wash words lists + for i in words: + word = "" + if isinstance(i, str): + word = i + else: + word = i.decode("utf-8") + + if word in ["", "", "", ""]: + continue + else: + middle_lists.append(word) + + # all alpha characters + for i, ch in enumerate(middle_lists): + word = "" + if "\u2581" in ch and i == 0: + word_item = "" + word = ch.replace("\u2581", "") + word_item += word + elif "\u2581" in ch and i != 0: + word_lists.append(word_item) + word_lists.append(" ") + word_item = "" + word = ch.replace("\u2581", "") + word_item += word + else: + word_item += ch + if word_item is not None: + word_lists.append(word_item) + # word_lists = abbr_dispose(word_lists) + real_word_lists = [] + for ch in word_lists: + if ch != " ": + if ch == "i": + ch = ch.replace("i", "I") + elif ch == "i'm": + ch = ch.replace("i'm", "I'm") + elif ch == "i've": + ch = ch.replace("i've", "I've") + elif ch == "i'll": + ch = ch.replace("i'll", "I'll") + real_word_lists.append(ch) + sentence = "".join(word_lists) + return sentence, real_word_lists + + +emo_dict = { + "<|HAPPY|>": "😊", + "<|SAD|>": "😔", + "<|ANGRY|>": "😡", + "<|NEUTRAL|>": "", + "<|FEARFUL|>": "😰", + "<|DISGUSTED|>": "🤢", + "<|SURPRISED|>": "😮", +} + +event_dict = { + "<|BGM|>": "🎼", + "<|Speech|>": "", + "<|Applause|>": "👏", + "<|Laughter|>": "😀", + "<|Cry|>": "😭", + "<|Sneeze|>": "🤧", + "<|Breath|>": "", + "<|Cough|>": "🤧", +} + +lang_dict = { + "<|zh|>": "<|lang|>", + "<|en|>": "<|lang|>", + "<|yue|>": "<|lang|>", + "<|ja|>": "<|lang|>", + "<|ko|>": "<|lang|>", + "<|nospeech|>": "<|lang|>", +} + +emoji_dict = { + "<|nospeech|><|Event_UNK|>": "❓", + "<|zh|>": "", + "<|en|>": "", + "<|yue|>": "", + "<|ja|>": "", + "<|ko|>": "", + "<|nospeech|>": "", + "<|HAPPY|>": "😊", + "<|SAD|>": "😔", + "<|ANGRY|>": "😡", + "<|NEUTRAL|>": "", + "<|BGM|>": "🎼", + "<|Speech|>": "", + "<|Applause|>": "👏", + "<|Laughter|>": "😀", + "<|FEARFUL|>": "😰", + "<|DISGUSTED|>": "🤢", + "<|SURPRISED|>": "😮", + "<|Cry|>": "😭", + "<|EMO_UNKNOWN|>": "", + "<|Sneeze|>": "🤧", + "<|Breath|>": "", + "<|Cough|>": "😷", + "<|Sing|>": "", + "<|Speech_Noise|>": "", + "<|withitn|>": "", + "<|woitn|>": "", + "<|GBG|>": "", + "<|Event_UNK|>": "", +} + +emo_set = {"😊", "😔", "😡", "😰", "🤢", "😮"} +event_set = { + "🎼", + "👏", + "😀", + "😭", + "🤧", + "😷", +} + + +def format_str_v2(s): + sptk_dict = {} + for sptk in emoji_dict: + sptk_dict[sptk] = s.count(sptk) + s = s.replace(sptk, "") + emo = "<|NEUTRAL|>" + for e in emo_dict: + if sptk_dict[e] > sptk_dict[emo]: + emo = e + for e in event_dict: + if sptk_dict[e] > 0: + s = event_dict[e] + s + s = s + emo_dict[emo] + + for emoji in emo_set.union(event_set): + s = s.replace(" " + emoji, emoji) + s = s.replace(emoji + " ", emoji) + return s.strip() + + +def rich_transcription_postprocess(s): + def get_emo(s): + return s[-1] if s[-1] in emo_set else None + + def get_event(s): + return s[0] if s[0] in event_set else None + + s = s.replace("<|nospeech|><|Event_UNK|>", "❓") + for lang in lang_dict: + s = s.replace(lang, "<|lang|>") + s_list = [format_str_v2(s_i).strip(" ") for s_i in s.split("<|lang|>")] + new_s = " " + s_list[0] + cur_ent_event = get_event(new_s) + for i in range(1, len(s_list)): + if len(s_list[i]) == 0: + continue + if get_event(s_list[i]) == cur_ent_event and get_event(s_list[i]) != None: + s_list[i] = s_list[i][1:] + # else: + cur_ent_event = get_event(s_list[i]) + if get_emo(s_list[i]) != None and get_emo(s_list[i]) == get_emo(new_s): + new_s = new_s[:-1] + new_s += s_list[i].strip().lstrip() + new_s = new_s.replace("The.", " ") + return new_s.strip() diff --git a/funasr/utils/run_bmodel.py b/funasr/utils/run_bmodel.py new file mode 100644 index 0000000000000000000000000000000000000000..1e9f407021afebe520ef184fb291b7b6064473cf --- /dev/null +++ b/funasr/utils/run_bmodel.py @@ -0,0 +1,35 @@ +# from wangyang +from tpu_perf.infer import SGInfer +import numpy as np +import time +#import torch +import os +class EngineOV: + + def __init__(self, model_path="", batch=1, device_id=0) : + # 如果环境变量中没有设置device_id，则使用默认值 + if "DEVICE_ID" in os.environ: + device_id = int(os.environ["DEVICE_ID"]) + print(">>>> device_id is in os.environ. and device_id = ",device_id) + self.model = SGInfer(model_path , batch=batch, devices=[device_id]) + + def __str__(self): + return "EngineOV: model_path={}, device_id={}".format(self.model_path,self.device_id) + + + def __call__(self, args): + start = time.time() + if isinstance(args, list): + values = args + elif isinstance(args, dict): + values = list(args.values()) + else: + raise TypeError("args is not list or dict") + # print(values) + #print(time.time() - start) + #start = time.time() + task_id = self.model.put(*values) + #print("put time : ",time.time() - start) + task_id, results, valid = self.model.get() + #print("get time : ",time.time() - start) + return results diff --git a/funasr/utils/speaker_utils.py b/funasr/utils/speaker_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e470849f8db48de8009b08b42cec0bca22dfedeb --- /dev/null +++ b/funasr/utils/speaker_utils.py @@ -0,0 +1,197 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +""" Some implementations are adapted from https://github.com/yuyq96/D-TDNN +""" + +import io +from typing import Union + +import librosa as sf +import numpy as np +import torch +import torch.nn.functional as F +import torchaudio.compliance.kaldi as Kaldi +from torch import nn + +from funasr.utils.modelscope_file import File + + +def check_audio_list(audio: list): + audio_dur = 0 + for i in range(len(audio)): + seg = audio[i] + assert seg[1] >= seg[0], "modelscope error: Wrong time stamps." + assert isinstance(seg[2], np.ndarray), "modelscope error: Wrong data type." + assert ( + int(seg[1] * 16000) - int(seg[0] * 16000) == seg[2].shape[0] + ), "modelscope error: audio data in list is inconsistent with time length." + if i > 0: + assert seg[0] >= audio[i - 1][1], "modelscope error: Wrong time stamps." + audio_dur += seg[1] - seg[0] + return audio_dur + # assert audio_dur > 5, 'modelscope error: The effective audio duration is too short.' + + +def sv_preprocess(inputs: Union[np.ndarray, list]): + output = [] + for i in range(len(inputs)): + if isinstance(inputs[i], str): + file_bytes = File.read(inputs[i]) + data, fs = sf.load(io.BytesIO(file_bytes), dtype="float32") + if len(data.shape) == 2: + data = data[:, 0] + data = torch.from_numpy(data).unsqueeze(0) + data = data.squeeze(0) + elif isinstance(inputs[i], np.ndarray): + assert len(inputs[i].shape) == 1, "modelscope error: Input array should be [N, T]" + data = inputs[i] + if data.dtype in ["int16", "int32", "int64"]: + data = (data / (1 << 15)).astype("float32") + else: + data = data.astype("float32") + data = torch.from_numpy(data) + else: + raise ValueError( + "modelscope error: The input type is restricted to audio address and nump array." + ) + output.append(data) + return output + + +def sv_chunk(vad_segments: list, fs=16000) -> list: + config = { + "seg_dur": 1.5, + "seg_shift": 0.75, + } + + def seg_chunk(seg_data): + seg_st = seg_data[0] + data = seg_data[2] + chunk_len = int(config["seg_dur"] * fs) + chunk_shift = int(config["seg_shift"] * fs) + last_chunk_ed = 0 + seg_res = [] + for chunk_st in range(0, data.shape[0], chunk_shift): + chunk_ed = min(chunk_st + chunk_len, data.shape[0]) + if chunk_ed <= last_chunk_ed: + break + last_chunk_ed = chunk_ed + chunk_st = max(0, chunk_ed - chunk_len) + chunk_data = data[chunk_st:chunk_ed] + if chunk_data.shape[0] < chunk_len: + chunk_data = np.pad(chunk_data, (0, chunk_len - chunk_data.shape[0]), "constant") + seg_res.append([chunk_st / fs + seg_st, chunk_ed / fs + seg_st, chunk_data]) + return seg_res + + segs = [] + for i, s in enumerate(vad_segments): + segs.extend(seg_chunk(s)) + + return segs + + +def extract_feature(audio): + features = [] + for au in audio: + feature = Kaldi.fbank(au.unsqueeze(0), num_mel_bins=80) + feature = feature - feature.mean(dim=0, keepdim=True) + features.append(feature.unsqueeze(0)) + features = torch.cat(features) + return features + + +def postprocess( + segments: list, vad_segments: list, labels: np.ndarray, embeddings: np.ndarray +) -> list: + assert len(segments) == len(labels) + labels = correct_labels(labels) + distribute_res = [] + for i in range(len(segments)): + distribute_res.append([segments[i][0], segments[i][1], labels[i]]) + # merge the same speakers chronologically + distribute_res = merge_seque(distribute_res) + + # accquire speaker center + spk_embs = [] + for i in range(labels.max() + 1): + spk_emb = embeddings[labels == i].mean(0) + spk_embs.append(spk_emb) + spk_embs = np.stack(spk_embs) + + def is_overlapped(t1, t2): + if t1 > t2 + 1e-4: + return True + return False + + # distribute the overlap region + for i in range(1, len(distribute_res)): + if is_overlapped(distribute_res[i - 1][1], distribute_res[i][0]): + p = (distribute_res[i][0] + distribute_res[i - 1][1]) / 2 + distribute_res[i][0] = p + distribute_res[i - 1][1] = p + + # smooth the result + distribute_res = smooth(distribute_res) + + return distribute_res + + +def correct_labels(labels): + labels_id = 0 + id2id = {} + new_labels = [] + for i in labels: + if i not in id2id: + id2id[i] = labels_id + labels_id += 1 + new_labels.append(id2id[i]) + return np.array(new_labels) + + +def merge_seque(distribute_res): + res = [distribute_res[0]] + for i in range(1, len(distribute_res)): + if distribute_res[i][2] != res[-1][2] or distribute_res[i][0] > res[-1][1]: + res.append(distribute_res[i]) + else: + res[-1][1] = distribute_res[i][1] + return res + + +def smooth(res, mindur=1): + # short segments are assigned to nearest speakers. + for i in range(len(res)): + res[i][0] = round(res[i][0], 2) + res[i][1] = round(res[i][1], 2) + if res[i][1] - res[i][0] < mindur: + if i == 0: + res[i][2] = res[i + 1][2] + elif i == len(res) - 1: + res[i][2] = res[i - 1][2] + elif res[i][0] - res[i - 1][1] <= res[i + 1][0] - res[i][1]: + res[i][2] = res[i - 1][2] + else: + res[i][2] = res[i + 1][2] + # merge the speakers + res = merge_seque(res) + + return res + + +def distribute_spk(sentence_list, sd_time_list): + sd_sentence_list = [] + for d in sentence_list: + sentence_start = d["ts_list"][0][0] + sentence_end = d["ts_list"][-1][1] + sentence_spk = 0 + max_overlap = 0 + for sd_time in sd_time_list: + spk_st, spk_ed, spk = sd_time + spk_st = spk_st * 1000 + spk_ed = spk_ed * 1000 + overlap = max(min(sentence_end, spk_ed) - max(sentence_start, spk_st), 0) + if overlap > max_overlap: + max_overlap = overlap + sentence_spk = spk + d["spk"] = sentence_spk + sd_sentence_list.append(d) + return sd_sentence_list diff --git a/funasr/utils/timestamp_tools.py b/funasr/utils/timestamp_tools.py new file mode 100644 index 0000000000000000000000000000000000000000..37ce8868331039cfc14f43fc5bcaf1d802e4d4f8 --- /dev/null +++ b/funasr/utils/timestamp_tools.py @@ -0,0 +1,278 @@ +import torch +import codecs +import logging +import argparse +import numpy as np + +# import edit_distance +from itertools import zip_longest + + +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 + + +def ts_prediction_lfr6_standard( + us_alphas, us_peaks, char_list, vad_offset=0.0, force_time_shift=-1.5, sil_in_str=True, upsample_rate=3, +): + if not len(char_list): + return "", [] + START_END_THRESHOLD = 5 + MAX_TOKEN_DURATION = 12 # 3 times upsampled + TIME_RATE=10.0 * 6 / 1000 / upsample_rate + if len(us_alphas.shape) == 2: + alphas, peaks = us_alphas[0], us_peaks[0] # support inference batch_size=1 only + else: + alphas, peaks = us_alphas, us_peaks + if char_list[-1] == "": + char_list = char_list[:-1] + fire_place = ( + torch.where(peaks >= 1.0 - 1e-4)[0].cpu().numpy() + force_time_shift + ) # total offset + if len(fire_place) != len(char_list) + 1: + alphas /= alphas.sum() / (len(char_list) + 1) + alphas = alphas.unsqueeze(0) + peaks = cif_wo_hidden(alphas, threshold=1.0 - 1e-4)[0] + fire_place = ( + torch.where(peaks >= 1.0 - 1e-4)[0].cpu().numpy() + force_time_shift + ) # total offset + num_frames = peaks.shape[0] + timestamp_list = [] + new_char_list = [] + # for bicif model trained with large data, cif2 actually fires when a character starts + # so treat the frames between two peaks as the duration of the former token + # fire_place = torch.where(peaks>=1.0-1e-4)[0].cpu().numpy() + force_time_shift # total offset + # assert num_peak == len(char_list) + 1 # number of peaks is supposed to be number of tokens + 1 + # begin silence + if fire_place[0] > START_END_THRESHOLD: + # char_list.insert(0, '') + timestamp_list.append([0.0, fire_place[0] * TIME_RATE]) + new_char_list.append("") + # tokens timestamp + for i in range(len(fire_place) - 1): + new_char_list.append(char_list[i]) + if MAX_TOKEN_DURATION < 0 or fire_place[i + 1] - fire_place[i] <= MAX_TOKEN_DURATION: + timestamp_list.append([fire_place[i] * TIME_RATE, fire_place[i + 1] * TIME_RATE]) + else: + # cut the duration to token and sil of the 0-weight frames last long + _split = fire_place[i] + MAX_TOKEN_DURATION + timestamp_list.append([fire_place[i] * TIME_RATE, _split * TIME_RATE]) + timestamp_list.append([_split * TIME_RATE, fire_place[i + 1] * TIME_RATE]) + new_char_list.append("") + # tail token and end silence + # new_char_list.append(char_list[-1]) + if num_frames - fire_place[-1] > START_END_THRESHOLD: + _end = (num_frames + fire_place[-1]) * 0.5 + # _end = fire_place[-1] + timestamp_list[-1][1] = _end * TIME_RATE + timestamp_list.append([_end * TIME_RATE, num_frames * TIME_RATE]) + new_char_list.append("") + else: + if len(timestamp_list)>0: + timestamp_list[-1][1] = num_frames * TIME_RATE + if vad_offset: # add offset time in model with vad + for i in range(len(timestamp_list)): + timestamp_list[i][0] = timestamp_list[i][0] + vad_offset / 1000.0 + timestamp_list[i][1] = timestamp_list[i][1] + vad_offset / 1000.0 + res_txt = "" + for char, timestamp in zip(new_char_list, timestamp_list): + # if char != '': + if not sil_in_str and char == "": + continue + res_txt += "{} {} {};".format( + char, str(timestamp[0] + 0.0005)[:5], str(timestamp[1] + 0.0005)[:5] + ) + res = [] + for char, timestamp in zip(new_char_list, timestamp_list): + if char != "": + res.append([int(timestamp[0] * 1000), int(timestamp[1] * 1000)]) + return res_txt, res + + +def timestamp_sentence( + punc_id_list, timestamp_postprocessed, text_postprocessed, return_raw_text=False +): + punc_list = ["，", "。", "？", "、"] + res = [] + if text_postprocessed is None: + return res + if timestamp_postprocessed is None: + return res + if len(timestamp_postprocessed) == 0: + return res + if len(text_postprocessed) == 0: + return res + + if punc_id_list is None or len(punc_id_list) == 0: + res.append( + { + "text": text_postprocessed.split(), + "start": timestamp_postprocessed[0][0], + "end": timestamp_postprocessed[-1][1], + "timestamp": timestamp_postprocessed, + } + ) + return res + if len(punc_id_list) != len(timestamp_postprocessed): + logging.warning("length mismatch between punc and timestamp") + sentence_text = "" + sentence_text_seg = "" + ts_list = [] + sentence_start = timestamp_postprocessed[0][0] + sentence_end = timestamp_postprocessed[0][1] + texts = text_postprocessed.split() + punc_stamp_text_list = list( + zip_longest(punc_id_list, timestamp_postprocessed, texts, fillvalue=None) + ) + for punc_stamp_text in punc_stamp_text_list: + punc_id, timestamp, text = punc_stamp_text + if sentence_start is None and timestamp is not None: + sentence_start = timestamp[0] + # sentence_text += text if text is not None else '' + if text is not None: + if "a" <= text[0] <= "z" or "A" <= text[0] <= "Z": + sentence_text += " " + text + elif len(sentence_text) and ( + "a" <= sentence_text[-1] <= "z" or "A" <= sentence_text[-1] <= "Z" + ): + sentence_text += " " + text + else: + sentence_text += text + sentence_text_seg += text + " " + ts_list.append(timestamp) + + punc_id = int(punc_id) if punc_id is not None else 1 + sentence_end = timestamp[1] if timestamp is not None else sentence_end + sentence_text_seg = ( + sentence_text_seg[:-1] if sentence_text_seg and sentence_text_seg[-1] == " " else sentence_text_seg + ) + if punc_id > 1: + sentence_text += punc_list[punc_id - 2] + if return_raw_text: + res.append( + { + "text": sentence_text, + "start": sentence_start, + "end": sentence_end, + "timestamp": ts_list, + "raw_text": sentence_text_seg, + } + ) + else: + res.append( + { + "text": sentence_text, + "start": sentence_start, + "end": sentence_end, + "timestamp": ts_list, + } + ) + sentence_text = "" + sentence_text_seg = "" + ts_list = [] + sentence_start = None + return res + + +def timestamp_sentence_en( + punc_id_list, timestamp_postprocessed, text_postprocessed, return_raw_text=False +): + punc_list = [",", ".", "?", ","] + res = [] + if text_postprocessed is None: + return res + if timestamp_postprocessed is None: + return res + if len(timestamp_postprocessed) == 0: + return res + if len(text_postprocessed) == 0: + return res + + if punc_id_list is None or len(punc_id_list) == 0: + res.append( + { + "text": text_postprocessed.split(), + "start": timestamp_postprocessed[0][0], + "end": timestamp_postprocessed[-1][1], + "timestamp": timestamp_postprocessed, + } + ) + return res + if len(punc_id_list) != len(timestamp_postprocessed): + logging.warning("length mismatch between punc and timestamp") + sentence_text = "" + sentence_text_seg = "" + ts_list = [] + sentence_start = timestamp_postprocessed[0][0] + sentence_end = timestamp_postprocessed[0][1] + texts = text_postprocessed.split() + punc_stamp_text_list = list( + zip_longest(punc_id_list, timestamp_postprocessed, texts, fillvalue=None) + ) + is_sentence_start = True + for punc_stamp_text in punc_stamp_text_list: + punc_id, timestamp, text = punc_stamp_text + # sentence_text += text if text is not None else '' + if text is not None: + if "a" <= text[0] <= "z" or "A" <= text[0] <= "Z": + sentence_text += " " + text + elif len(sentence_text) and ( + "a" <= sentence_text[-1] <= "z" or "A" <= sentence_text[-1] <= "Z" + ): + sentence_text += " " + text + else: + sentence_text += text + sentence_text_seg += text + " " + ts_list.append(timestamp) + + punc_id = int(punc_id) if punc_id is not None else 1 + sentence_end = timestamp[1] if timestamp is not None else sentence_end + sentence_text = sentence_text[1:] if sentence_text[0] == ' ' else sentence_text + if is_sentence_start: + sentence_start = timestamp[0] if timestamp is not None else sentence_start + is_sentence_start = False + if punc_id > 1: + is_sentence_start = True + sentence_text += punc_list[punc_id - 2] + sentence_text_seg = ( + sentence_text_seg[:-1] if sentence_text_seg[-1] == " " else sentence_text_seg + ) + if return_raw_text: + res.append( + { + "text": sentence_text, + "start": sentence_start, + "end": sentence_end, + "timestamp": ts_list, + "raw_text": sentence_text_seg, + } + ) + else: + res.append( + { + "text": sentence_text, + "start": sentence_start, + "end": sentence_end, + "timestamp": ts_list, + } + ) + sentence_text = "" + sentence_text_seg = "" + ts_list = [] + return res diff --git a/funasr/utils/torch_function.py b/funasr/utils/torch_function.py new file mode 100644 index 0000000000000000000000000000000000000000..f637bbf82e5848dd4ae75133dd457eca46429ab2 --- /dev/null +++ b/funasr/utils/torch_function.py @@ -0,0 +1,84 @@ +from typing import Optional + +import torch +import torch.nn as nn + +import numpy as np + + +class MakePadMask(nn.Module): + def __init__(self, max_seq_len=512, flip=True): + super().__init__() + if flip: + self.mask_pad = torch.Tensor(1 - np.tri(max_seq_len)).type(torch.bool) + else: + self.mask_pad = torch.Tensor(np.tri(max_seq_len)).type(torch.bool) + + def forward(self, lengths, xs=None, length_dim=-1, maxlen=None): + """Make mask tensor containing indices of padded part. + This implementation creates the same mask tensor with original make_pad_mask, + which can be converted into onnx format. + Dimension length of xs should be 2 or 3. + """ + if length_dim == 0: + raise ValueError("length_dim cannot be 0: {}".format(length_dim)) + + if xs is not None and len(xs.shape) == 3: + if length_dim == 1: + lengths = lengths.unsqueeze(1).expand(*xs.transpose(1, 2).shape[:2]) + else: + lengths = lengths.unsqueeze(1).expand(*xs.shape[:2]) + + if maxlen is not None: + m = maxlen + elif xs is not None: + m = xs.shape[-1] + else: + m = torch.max(lengths) + + mask = self.mask_pad[lengths - 1][..., :m].type(torch.float32) + + if length_dim == 1: + return mask.transpose(1, 2) + else: + return mask + + +class sequence_mask(nn.Module): + def __init__(self, max_seq_len=512, flip=True): + super().__init__() + + def forward(self, lengths, max_seq_len=None, dtype=torch.float32, device=None): + if max_seq_len is None: + max_seq_len = lengths.max() + row_vector = torch.arange(0, max_seq_len, 1).to(lengths.device) + matrix = torch.unsqueeze(lengths, dim=-1) + mask = row_vector < matrix + + return mask.type(dtype).to(device) if device is not None else mask.type(dtype) + + +def normalize( + input: torch.Tensor, p: float = 2.0, dim: int = 1, out: Optional[torch.Tensor] = None +) -> torch.Tensor: + if out is None: + denom = input.norm(p, dim, keepdim=True).expand_as(input) + return input / denom + else: + denom = input.norm(p, dim, keepdim=True).expand_as(input) + return torch.div(input, denom, out=out) + + +def subsequent_mask(size: torch.Tensor): + return torch.ones(size, size).tril() + + +def MakePadMask_test(): + feats_length = torch.tensor([10]).type(torch.long) + mask_fn = MakePadMask() + mask = mask_fn(feats_length) + print(mask) + + +if __name__ == "__main__": + MakePadMask_test() diff --git a/funasr/utils/type_utils.py b/funasr/utils/type_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..6b36f9c4b87ed9258a5d1e254ba298ed5dbc01d2 --- /dev/null +++ b/funasr/utils/type_utils.py @@ -0,0 +1,149 @@ +from distutils.util import strtobool +from typing import Optional +from typing import Tuple +from typing import Union + +import humanfriendly + + +def str2bool(value: str) -> bool: + return bool(strtobool(value)) + + +def remove_parenthesis(value: str): + value = value.strip() + if value.startswith("(") and value.endswith(")"): + value = value[1:-1] + elif value.startswith("[") and value.endswith("]"): + value = value[1:-1] + return value + + +def remove_quotes(value: str): + value = value.strip() + if value.startswith('"') and value.endswith('"'): + value = value[1:-1] + elif value.startswith("'") and value.endswith("'"): + value = value[1:-1] + return value + + +def int_or_none(value: str) -> Optional[int]: + """int_or_none. + + Examples: + >>> import argparse + >>> parser = argparse.ArgumentParser() + >>> _ = parser.add_argument('--foo', type=int_or_none) + >>> parser.parse_args(['--foo', '456']) + Namespace(foo=456) + >>> parser.parse_args(['--foo', 'none']) + Namespace(foo=None) + >>> parser.parse_args(['--foo', 'null']) + Namespace(foo=None) + >>> parser.parse_args(['--foo', 'nil']) + Namespace(foo=None) + + """ + if value.strip().lower() in ("none", "null", "nil"): + return None + return int(value) + + +def float_or_none(value: str) -> Optional[float]: + """float_or_none. + + Examples: + >>> import argparse + >>> parser = argparse.ArgumentParser() + >>> _ = parser.add_argument('--foo', type=float_or_none) + >>> parser.parse_args(['--foo', '4.5']) + Namespace(foo=4.5) + >>> parser.parse_args(['--foo', 'none']) + Namespace(foo=None) + >>> parser.parse_args(['--foo', 'null']) + Namespace(foo=None) + >>> parser.parse_args(['--foo', 'nil']) + Namespace(foo=None) + + """ + if value.strip().lower() in ("none", "null", "nil"): + return None + return float(value) + + +def humanfriendly_parse_size_or_none(value) -> Optional[float]: + if value.strip().lower() in ("none", "null", "nil"): + return None + return humanfriendly.parse_size(value) + + +def str_or_int(value: str) -> Union[str, int]: + try: + return int(value) + except ValueError: + return value + + +def str_or_none(value: str) -> Optional[str]: + """str_or_none. + + Examples: + >>> import argparse + >>> parser = argparse.ArgumentParser() + >>> _ = parser.add_argument('--foo', type=str_or_none) + >>> parser.parse_args(['--foo', 'aaa']) + Namespace(foo='aaa') + >>> parser.parse_args(['--foo', 'none']) + Namespace(foo=None) + >>> parser.parse_args(['--foo', 'null']) + Namespace(foo=None) + >>> parser.parse_args(['--foo', 'nil']) + Namespace(foo=None) + + """ + if value.strip().lower() in ("none", "null", "nil"): + return None + return value + + +def str2pair_str(value: str) -> Tuple[str, str]: + """str2pair_str. + + Examples: + >>> import argparse + >>> str2pair_str('abc,def ') + ('abc', 'def') + >>> parser = argparse.ArgumentParser() + >>> _ = parser.add_argument('--foo', type=str2pair_str) + >>> parser.parse_args(['--foo', 'abc,def']) + Namespace(foo=('abc', 'def')) + + """ + value = remove_parenthesis(value) + a, b = value.split(",") + + # Workaround for configargparse issues: + # If the list values are given from yaml file, + # the value givent to type() is shaped as python-list, + # e.g. ['a', 'b', 'c'], + # so we need to remove double quotes from it. + return remove_quotes(a), remove_quotes(b) + + +def str2triple_str(value: str) -> Tuple[str, str, str]: + """str2triple_str. + + Examples: + >>> str2triple_str('abc,def ,ghi') + ('abc', 'def', 'ghi') + """ + value = remove_parenthesis(value) + a, b, c = value.split(",") + + # Workaround for configargparse issues: + # If the list values are given from yaml file, + # the value givent to type() is shaped as python-list, + # e.g. ['a', 'b', 'c'], + # so we need to remove quotes from it. + return remove_quotes(a), remove_quotes(b), remove_quotes(c) diff --git a/funasr/utils/vad_utils.py b/funasr/utils/vad_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f87a85e711a8fe981f8125eae691bc820dd2a67a --- /dev/null +++ b/funasr/utils/vad_utils.py @@ -0,0 +1,59 @@ +import torch +from torch.nn.utils.rnn import pad_sequence + + +def slice_padding_fbank(speech, speech_lengths, vad_segments): + speech_list = [] + speech_lengths_list = [] + for i, segment in enumerate(vad_segments): + + bed_idx = int(segment[0][0] * 16) + end_idx = min(int(segment[0][1] * 16), speech_lengths[0]) + speech_i = speech[0, bed_idx:end_idx] + speech_lengths_i = end_idx - bed_idx + speech_list.append(speech_i) + speech_lengths_list.append(speech_lengths_i) + feats_pad = pad_sequence(speech_list, batch_first=True, padding_value=0.0) + speech_lengths_pad = torch.Tensor(speech_lengths_list).int() + return feats_pad, speech_lengths_pad + + +def slice_padding_audio_samples(speech, speech_lengths, vad_segments): + speech_list = [] + speech_lengths_list = [] + for i, segment in enumerate(vad_segments): + bed_idx = int(segment[0][0] * 16) + end_idx = min(int(segment[0][1] * 16), speech_lengths) + speech_i = speech[bed_idx:end_idx] + speech_lengths_i = end_idx - bed_idx + speech_list.append(speech_i) + speech_lengths_list.append(speech_lengths_i) + + return speech_list, speech_lengths_list + + +def merge_vad(vad_result, max_length=15000, min_length=0): + new_result = [] + if len(vad_result) <= 1: + return vad_result + time_step = [t[0] for t in vad_result] + [t[1] for t in vad_result] + time_step = sorted(list(set(time_step))) + if len(time_step) == 0: + return [] + bg = 0 + for i in range(len(time_step) - 1): + time = time_step[i] + if time_step[i + 1] - bg < max_length: + continue + if time - bg > min_length: + new_result.append([bg, time]) + # if time - bg < max_length * 1.5: + # new_result.append([bg, time]) + # else: + # split_num = int(time - bg) // max_length + 1 + # spl_l = int(time - bg) // split_num + # for j in range(split_num): + # new_result.append([bg + j * spl_l, bg + (j + 1) * spl_l]) + bg = time + new_result.append([bg, time_step[-1]]) + return new_result \ No newline at end of file diff --git a/funasr/utils/version_checker.py b/funasr/utils/version_checker.py new file mode 100644 index 0000000000000000000000000000000000000000..9bc11deb9ecd4bf4237a12703132011657c5f5be --- /dev/null +++ b/funasr/utils/version_checker.py @@ -0,0 +1,32 @@ +from packaging import version +from funasr import __version__ # Ensure that __version__ is defined in your package's __init__.py + + +def get_pypi_version(package_name): + import requests + + url = f"https://pypi.org/pypi/{package_name}/json" + response = requests.get(url) + if response.status_code == 200: + data = response.json() + return version.parse(data["info"]["version"]) + else: + raise Exception("Failed to retrieve version information from PyPI.") + + +def check_for_update(disable=False): + current_version = version.parse(__version__) + print(f"funasr version: {current_version}.") + if disable: + return + + print( + "Check update of funasr, and it would cost few times. You may disable it by set `disable_update=True` in AutoModel" + ) + pypi_version = get_pypi_version("funasr") + + if current_version < pypi_version: + print(f"New version is available: {pypi_version}.") + print('Please use the command "pip install -U funasr" to upgrade.') + else: + print(f"You are using the latest version of funasr-{current_version}") diff --git a/funasr/version.txt b/funasr/version.txt new file mode 100644 index 0000000000000000000000000000000000000000..7e099ec5d75abf22ee60bb93ebb1803ec9b2687a --- /dev/null +++ b/funasr/version.txt @@ -0,0 +1 @@ +1.2.6 \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..ac6b0dca9568d3e9b260d2c0ee21c23b187d5b2c --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,39 @@ +[project] +name = "funasr-bmodel" +version = "0.1.0" +description = "Add your description here" +readme = "README.md" +requires-python = ">=3.8" +dependencies = [ + "hydra-core>=1.3.2", + "kaldiio>=2.18.1", + "librosa>=0.11.0", + "omegaconf>=2.3.0", + "rotary-embedding-torch>=0.8.9", + "six>=1.17.0", + "torch-complex>=0.4.4", + "torchaudio>=2.5.1", + "tpu-perf", + "tqdm>=4.67.1", + "umap-learn>=0.5.7", +] + + +[tool.uv.sources] +torch = { index = "pytorch" } +torchvision = { index = "pytorch" } +torchaudio = { index = "pytorch" } +tpu-perf = { path = "bmodel/tpu_perf-1.2.35-py3-none-manylinux2014_aarch64.whl" } + +[tool.uv] +python-downloads = "automatic" +python-install-mirror = "https://ghfast.top/github.com/indygreg/python-build-standalone/releases/download" + +[[tool.uv.index]] +url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" +default = true + +[[tool.uv.index]] +name = "pytorch" +url = "https://download.pytorch.org/whl/cpu" +explicit = true diff --git a/system_phrases.txt b/system_phrases.txt new file mode 100644 index 0000000000000000000000000000000000000000..16da10ab1eb7c1b04190ccde8c85615593bf27c3 --- /dev/null +++ b/system_phrases.txt @@ -0,0 +1,10 @@ +在吴照明的道路行驶 +夜间通过急弯、坡路、拱桥、人行道或没有交通信号灯路口 +夜间通过急湾坡路拱桥神行道或没有交通信号 +夜间通过吉湾坡路 +拱桥神频道或没有交通信号灯路口 +或没有交通信号灯路口 +桥神频道或没有交通信号灯路口 +再有信号灯控制路口转弯 +在照明不良的道路行驶 +在照明良好的道路上行驶 \ No newline at end of file diff --git a/system_ref.wav b/system_ref.wav new file mode 100644 index 0000000000000000000000000000000000000000..5db2d462584406426c3a5cf03204d64b9990aeda --- /dev/null +++ b/system_ref.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:418b16c8611ad2ef3b8e254225442e25919ff599cf76558e3646237bd6101295 +size 217004 diff --git a/uv.lock b/uv.lock new file mode 100644 index 0000000000000000000000000000000000000000..540e6350835a0881b80f213687b860376c7cdc0b --- /dev/null +++ b/uv.lock @@ -0,0 +1,2975 @@ +version = 1 +revision = 3 +requires-python = ">=3.8" +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.13.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'darwin')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "python_full_version < '3.9' and sys_platform != 'darwin'", + "(python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'darwin')", + "(python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform == 'darwin')", + "(python_full_version == '3.10.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'darwin')", + "(python_full_version == '3.9.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform == 'darwin')", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] + +[[package]] +name = "antlr4-python3-runtime" +version = "4.9.3" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3e/38/7859ff46355f76f8d19459005ca000b6e7012f2f1ca597746cbcd1fbfe5e/antlr4-python3-runtime-4.9.3.tar.gz", hash = "sha256:f224469b4168294902bb1efa80a8bf7855f24c99aef99cbefc1bcd3cce77881b", size = 117034, upload-time = "2021-11-06T17:52:23.524Z" } + +[[package]] +name = "audioop-lts" +version = "0.2.2" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/38/53/946db57842a50b2da2e0c1e34bd37f36f5aadba1a929a3971c5d7841dbca/audioop_lts-0.2.2.tar.gz", hash = "sha256:64d0c62d88e67b98a1a5e71987b7aa7b5bcffc7dcee65b635823dbdd0a8dbbd0", size = 30686, upload-time = "2025-08-05T16:43:17.409Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/de/d4/94d277ca941de5a507b07f0b592f199c22454eeaec8f008a286b3fbbacd6/audioop_lts-0.2.2-cp313-abi3-macosx_10_13_universal2.whl", hash = "sha256:fd3d4602dc64914d462924a08c1a9816435a2155d74f325853c1f1ac3b2d9800", size = 46523, upload-time = "2025-08-05T16:42:20.836Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f8/5a/656d1c2da4b555920ce4177167bfeb8623d98765594af59702c8873f60ec/audioop_lts-0.2.2-cp313-abi3-macosx_10_13_x86_64.whl", hash = "sha256:550c114a8df0aafe9a05442a1162dfc8fec37e9af1d625ae6060fed6e756f303", size = 27455, upload-time = "2025-08-05T16:42:22.283Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1b/83/ea581e364ce7b0d41456fb79d6ee0ad482beda61faf0cab20cbd4c63a541/audioop_lts-0.2.2-cp313-abi3-macosx_11_0_arm64.whl", hash = "sha256:9a13dc409f2564de15dd68be65b462ba0dde01b19663720c68c1140c782d1d75", size = 26997, upload-time = "2025-08-05T16:42:23.849Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b8/3b/e8964210b5e216e5041593b7d33e97ee65967f17c282e8510d19c666dab4/audioop_lts-0.2.2-cp313-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:51c916108c56aa6e426ce611946f901badac950ee2ddaf302b7ed35d9958970d", size = 85844, upload-time = "2025-08-05T16:42:25.208Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c7/2e/0a1c52faf10d51def20531a59ce4c706cb7952323b11709e10de324d6493/audioop_lts-0.2.2-cp313-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:47eba38322370347b1c47024defbd36374a211e8dd5b0dcbce7b34fdb6f8847b", size = 85056, upload-time = "2025-08-05T16:42:26.559Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/75/e8/cd95eef479656cb75ab05dfece8c1f8c395d17a7c651d88f8e6e291a63ab/audioop_lts-0.2.2-cp313-abi3-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ba7c3a7e5f23e215cb271516197030c32aef2e754252c4c70a50aaff7031a2c8", size = 93892, upload-time = "2025-08-05T16:42:27.902Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5c/1e/a0c42570b74f83efa5cca34905b3eef03f7ab09fe5637015df538a7f3345/audioop_lts-0.2.2-cp313-abi3-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:def246fe9e180626731b26e89816e79aae2276f825420a07b4a647abaa84becc", size = 96660, upload-time = "2025-08-05T16:42:28.9Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/50/d5/8a0ae607ca07dbb34027bac8db805498ee7bfecc05fd2c148cc1ed7646e7/audioop_lts-0.2.2-cp313-abi3-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e160bf9df356d841bb6c180eeeea1834085464626dc1b68fa4e1d59070affdc3", size = 79143, upload-time = "2025-08-05T16:42:29.929Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/12/17/0d28c46179e7910bfb0bb62760ccb33edb5de973052cb2230b662c14ca2e/audioop_lts-0.2.2-cp313-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:4b4cd51a57b698b2d06cb9993b7ac8dfe89a3b2878e96bc7948e9f19ff51dba6", size = 84313, upload-time = "2025-08-05T16:42:30.949Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/84/ba/bd5d3806641564f2024e97ca98ea8f8811d4e01d9b9f9831474bc9e14f9e/audioop_lts-0.2.2-cp313-abi3-musllinux_1_2_ppc64le.whl", hash = "sha256:4a53aa7c16a60a6857e6b0b165261436396ef7293f8b5c9c828a3a203147ed4a", size = 93044, upload-time = "2025-08-05T16:42:31.959Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f9/5e/435ce8d5642f1f7679540d1e73c1c42d933331c0976eb397d1717d7f01a3/audioop_lts-0.2.2-cp313-abi3-musllinux_1_2_riscv64.whl", hash = "sha256:3fc38008969796f0f689f1453722a0f463da1b8a6fbee11987830bfbb664f623", size = 78766, upload-time = "2025-08-05T16:42:33.302Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ae/3b/b909e76b606cbfd53875693ec8c156e93e15a1366a012f0b7e4fb52d3c34/audioop_lts-0.2.2-cp313-abi3-musllinux_1_2_s390x.whl", hash = "sha256:15ab25dd3e620790f40e9ead897f91e79c0d3ce65fe193c8ed6c26cffdd24be7", size = 87640, upload-time = "2025-08-05T16:42:34.854Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/30/e7/8f1603b4572d79b775f2140d7952f200f5e6c62904585d08a01f0a70393a/audioop_lts-0.2.2-cp313-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:03f061a1915538fd96272bac9551841859dbb2e3bf73ebe4a23ef043766f5449", size = 86052, upload-time = "2025-08-05T16:42:35.839Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b5/96/c37846df657ccdda62ba1ae2b6534fa90e2e1b1742ca8dcf8ebd38c53801/audioop_lts-0.2.2-cp313-abi3-win32.whl", hash = "sha256:3bcddaaf6cc5935a300a8387c99f7a7fbbe212a11568ec6cf6e4bc458c048636", size = 26185, upload-time = "2025-08-05T16:42:37.04Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/34/a5/9d78fdb5b844a83da8a71226c7bdae7cc638861085fff7a1d707cb4823fa/audioop_lts-0.2.2-cp313-abi3-win_amd64.whl", hash = "sha256:a2c2a947fae7d1062ef08c4e369e0ba2086049a5e598fda41122535557012e9e", size = 30503, upload-time = "2025-08-05T16:42:38.427Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/34/25/20d8fde083123e90c61b51afb547bb0ea7e77bab50d98c0ab243d02a0e43/audioop_lts-0.2.2-cp313-abi3-win_arm64.whl", hash = "sha256:5f93a5db13927a37d2d09637ccca4b2b6b48c19cd9eda7b17a2e9f77edee6a6f", size = 24173, upload-time = "2025-08-05T16:42:39.704Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/58/a7/0a764f77b5c4ac58dc13c01a580f5d32ae8c74c92020b961556a43e26d02/audioop_lts-0.2.2-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:73f80bf4cd5d2ca7814da30a120de1f9408ee0619cc75da87d0641273d202a09", size = 47096, upload-time = "2025-08-05T16:42:40.684Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/aa/ed/ebebedde1a18848b085ad0fa54b66ceb95f1f94a3fc04f1cd1b5ccb0ed42/audioop_lts-0.2.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:106753a83a25ee4d6f473f2be6b0966fc1c9af7e0017192f5531a3e7463dce58", size = 27748, upload-time = "2025-08-05T16:42:41.992Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/cb/6e/11ca8c21af79f15dbb1c7f8017952ee8c810c438ce4e2b25638dfef2b02c/audioop_lts-0.2.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:fbdd522624141e40948ab3e8cdae6e04c748d78710e9f0f8d4dae2750831de19", size = 27329, upload-time = "2025-08-05T16:42:42.987Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/84/52/0022f93d56d85eec5da6b9da6a958a1ef09e80c39f2cc0a590c6af81dcbb/audioop_lts-0.2.2-cp313-cp313t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:143fad0311e8209ece30a8dbddab3b65ab419cbe8c0dde6e8828da25999be911", size = 92407, upload-time = "2025-08-05T16:42:44.336Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/87/1d/48a889855e67be8718adbc7a01f3c01d5743c325453a5e81cf3717664aad/audioop_lts-0.2.2-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:dfbbc74ec68a0fd08cfec1f4b5e8cca3d3cd7de5501b01c4b5d209995033cde9", size = 91811, upload-time = "2025-08-05T16:42:45.325Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/98/a6/94b7213190e8077547ffae75e13ed05edc488653c85aa5c41472c297d295/audioop_lts-0.2.2-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:cfcac6aa6f42397471e4943e0feb2244549db5c5d01efcd02725b96af417f3fe", size = 100470, upload-time = "2025-08-05T16:42:46.468Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e9/e9/78450d7cb921ede0cfc33426d3a8023a3bda755883c95c868ee36db8d48d/audioop_lts-0.2.2-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:752d76472d9804ac60f0078c79cdae8b956f293177acd2316cd1e15149aee132", size = 103878, upload-time = "2025-08-05T16:42:47.576Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4f/e2/cd5439aad4f3e34ae1ee852025dc6aa8f67a82b97641e390bf7bd9891d3e/audioop_lts-0.2.2-cp313-cp313t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:83c381767e2cc10e93e40281a04852facc4cd9334550e0f392f72d1c0a9c5753", size = 84867, upload-time = "2025-08-05T16:42:49.003Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/68/4b/9d853e9076c43ebba0d411e8d2aa19061083349ac695a7d082540bad64d0/audioop_lts-0.2.2-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:c0022283e9556e0f3643b7c3c03f05063ca72b3063291834cca43234f20c60bb", size = 90001, upload-time = "2025-08-05T16:42:50.038Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/58/26/4bae7f9d2f116ed5593989d0e521d679b0d583973d203384679323d8fa85/audioop_lts-0.2.2-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:a2d4f1513d63c795e82948e1305f31a6d530626e5f9f2605408b300ae6095093", size = 99046, upload-time = "2025-08-05T16:42:51.111Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b2/67/a9f4fb3e250dda9e9046f8866e9fa7d52664f8985e445c6b4ad6dfb55641/audioop_lts-0.2.2-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:c9c8e68d8b4a56fda8c025e538e639f8c5953f5073886b596c93ec9b620055e7", size = 84788, upload-time = "2025-08-05T16:42:52.198Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/70/f7/3de86562db0121956148bcb0fe5b506615e3bcf6e63c4357a612b910765a/audioop_lts-0.2.2-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:96f19de485a2925314f5020e85911fb447ff5fbef56e8c7c6927851b95533a1c", size = 94472, upload-time = "2025-08-05T16:42:53.59Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f1/32/fd772bf9078ae1001207d2df1eef3da05bea611a87dd0e8217989b2848fa/audioop_lts-0.2.2-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:e541c3ef484852ef36545f66209444c48b28661e864ccadb29daddb6a4b8e5f5", size = 92279, upload-time = "2025-08-05T16:42:54.632Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4f/41/affea7181592ab0ab560044632571a38edaf9130b84928177823fbf3176a/audioop_lts-0.2.2-cp313-cp313t-win32.whl", hash = "sha256:d5e73fa573e273e4f2e5ff96f9043858a5e9311e94ffefd88a3186a910c70917", size = 26568, upload-time = "2025-08-05T16:42:55.627Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/28/2b/0372842877016641db8fc54d5c88596b542eec2f8f6c20a36fb6612bf9ee/audioop_lts-0.2.2-cp313-cp313t-win_amd64.whl", hash = "sha256:9191d68659eda01e448188f60364c7763a7ca6653ed3f87ebb165822153a8547", size = 30942, upload-time = "2025-08-05T16:42:56.674Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ee/ca/baf2b9cc7e96c179bb4a54f30fcd83e6ecb340031bde68f486403f943768/audioop_lts-0.2.2-cp313-cp313t-win_arm64.whl", hash = "sha256:c174e322bb5783c099aaf87faeb240c8d210686b04bd61dfd05a8e5a83d88969", size = 24603, upload-time = "2025-08-05T16:42:57.571Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5c/73/413b5a2804091e2c7d5def1d618e4837f1cb82464e230f827226278556b7/audioop_lts-0.2.2-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:f9ee9b52f5f857fbaf9d605a360884f034c92c1c23021fb90b2e39b8e64bede6", size = 47104, upload-time = "2025-08-05T16:42:58.518Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ae/8c/daa3308dc6593944410c2c68306a5e217f5c05b70a12e70228e7dd42dc5c/audioop_lts-0.2.2-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:49ee1a41738a23e98d98b937a0638357a2477bc99e61b0f768a8f654f45d9b7a", size = 27754, upload-time = "2025-08-05T16:43:00.132Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4e/86/c2e0f627168fcf61781a8f72cab06b228fe1da4b9fa4ab39cfb791b5836b/audioop_lts-0.2.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:5b00be98ccd0fc123dcfad31d50030d25fcf31488cde9e61692029cd7394733b", size = 27332, upload-time = "2025-08-05T16:43:01.666Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c7/bd/35dce665255434f54e5307de39e31912a6f902d4572da7c37582809de14f/audioop_lts-0.2.2-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:a6d2e0f9f7a69403e388894d4ca5ada5c47230716a03f2847cfc7bd1ecb589d6", size = 92396, upload-time = "2025-08-05T16:43:02.991Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2d/d2/deeb9f51def1437b3afa35aeb729d577c04bcd89394cb56f9239a9f50b6f/audioop_lts-0.2.2-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f9b0b8a03ef474f56d1a842af1a2e01398b8f7654009823c6d9e0ecff4d5cfbf", size = 91811, upload-time = "2025-08-05T16:43:04.096Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/76/3b/09f8b35b227cee28cc8231e296a82759ed80c1a08e349811d69773c48426/audioop_lts-0.2.2-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:2b267b70747d82125f1a021506565bdc5609a2b24bcb4773c16d79d2bb260bbd", size = 100483, upload-time = "2025-08-05T16:43:05.085Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0b/15/05b48a935cf3b130c248bfdbdea71ce6437f5394ee8533e0edd7cfd93d5e/audioop_lts-0.2.2-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:0337d658f9b81f4cd0fdb1f47635070cc084871a3d4646d9de74fdf4e7c3d24a", size = 103885, upload-time = "2025-08-05T16:43:06.197Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/83/80/186b7fce6d35b68d3d739f228dc31d60b3412105854edb975aa155a58339/audioop_lts-0.2.2-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:167d3b62586faef8b6b2275c3218796b12621a60e43f7e9d5845d627b9c9b80e", size = 84899, upload-time = "2025-08-05T16:43:07.291Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/49/89/c78cc5ac6cb5828f17514fb12966e299c850bc885e80f8ad94e38d450886/audioop_lts-0.2.2-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:0d9385e96f9f6da847f4d571ce3cb15b5091140edf3db97276872647ce37efd7", size = 89998, upload-time = "2025-08-05T16:43:08.335Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4c/4b/6401888d0c010e586c2ca50fce4c903d70a6bb55928b16cfbdfd957a13da/audioop_lts-0.2.2-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:48159d96962674eccdca9a3df280e864e8ac75e40a577cc97c5c42667ffabfc5", size = 99046, upload-time = "2025-08-05T16:43:09.367Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/de/f8/c874ca9bb447dae0e2ef2e231f6c4c2b0c39e31ae684d2420b0f9e97ee68/audioop_lts-0.2.2-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:8fefe5868cd082db1186f2837d64cfbfa78b548ea0d0543e9b28935ccce81ce9", size = 84843, upload-time = "2025-08-05T16:43:10.749Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3e/c0/0323e66f3daebc13fd46b36b30c3be47e3fc4257eae44f1e77eb828c703f/audioop_lts-0.2.2-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:58cf54380c3884fb49fdd37dfb7a772632b6701d28edd3e2904743c5e1773602", size = 94490, upload-time = "2025-08-05T16:43:12.131Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/98/6b/acc7734ac02d95ab791c10c3f17ffa3584ccb9ac5c18fd771c638ed6d1f5/audioop_lts-0.2.2-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:088327f00488cdeed296edd9215ca159f3a5a5034741465789cad403fcf4bec0", size = 92297, upload-time = "2025-08-05T16:43:13.139Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/13/c3/c3dc3f564ce6877ecd2a05f8d751b9b27a8c320c2533a98b0c86349778d0/audioop_lts-0.2.2-cp314-cp314t-win32.whl", hash = "sha256:068aa17a38b4e0e7de771c62c60bbca2455924b67a8814f3b0dee92b5820c0b3", size = 27331, upload-time = "2025-08-05T16:43:14.19Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/72/bb/b4608537e9ffcb86449091939d52d24a055216a36a8bf66b936af8c3e7ac/audioop_lts-0.2.2-cp314-cp314t-win_amd64.whl", hash = "sha256:a5bf613e96f49712073de86f20dbdd4014ca18efd4d34ed18c75bd808337851b", size = 31697, upload-time = "2025-08-05T16:43:15.193Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f6/22/91616fe707a5c5510de2cac9b046a30defe7007ba8a0c04f9c08f27df312/audioop_lts-0.2.2-cp314-cp314t-win_arm64.whl", hash = "sha256:b492c3b040153e68b9fdaff5913305aaaba5bb433d8a7f73d5cf6a64ed3cc1dd", size = 25206, upload-time = "2025-08-05T16:43:16.444Z" }, +] + +[[package]] +name = "audioread" +version = "3.0.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/db/d2/87016ca9f083acadffb2d8da59bfa3253e4da7eeb9f71fb8e7708dc97ecd/audioread-3.0.1.tar.gz", hash = "sha256:ac5460a5498c48bdf2e8e767402583a4dcd13f4414d286f42ce4379e8b35066d" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/57/8d/30aa32745af16af0a9a650115fbe81bde7c610ed5c21b381fca0196f3a7f/audioread-3.0.1-py3-none-any.whl", hash = "sha256:4cdce70b8adc0da0a3c9e0d85fb10b3ace30fbdf8d1670fd443929b61d117c33" }, +] + +[[package]] +name = "certifi" +version = "2025.8.3" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/dc/67/960ebe6bf230a96cda2e0abcf73af550ec4f090005363542f0765df162e0/certifi-2025.8.3.tar.gz", hash = "sha256:e564105f78ded564e3ae7c923924435e1daa7463faeab5bb932bc53ffae63407", size = 162386, upload-time = "2025-08-03T03:07:47.08Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e5/48/1549795ba7742c948d2ad169c1c8cdbae65bc450d6cd753d124b17c8cd32/certifi-2025.8.3-py3-none-any.whl", hash = "sha256:f6c12493cfb1b06ba2ff328595af9350c65d6644968e5d3a2ffd78699af217a5", size = 161216, upload-time = "2025-08-03T03:07:45.777Z" }, +] + +[[package]] +name = "cffi" +version = "1.17.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +dependencies = [ + { name = "pycparser" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fc/97/c783634659c2920c3fc70419e3af40972dbaf758daa229a7d6ea6135c90d/cffi-1.17.1.tar.gz", hash = "sha256:1c39c6016c32bc48dd54561950ebd6836e1670f2ae46128f67cf49e789c52824", size = 516621, upload-time = "2024-09-04T20:45:21.852Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/90/07/f44ca684db4e4f08a3fdc6eeb9a0d15dc6883efc7b8c90357fdbf74e186c/cffi-1.17.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:df8b1c11f177bc2313ec4b2d46baec87a5f3e71fc8b45dab2ee7cae86d9aba14", size = 182191, upload-time = "2024-09-04T20:43:30.027Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/08/fd/cc2fedbd887223f9f5d170c96e57cbf655df9831a6546c1727ae13fa977a/cffi-1.17.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:8f2cdc858323644ab277e9bb925ad72ae0e67f69e804f4898c070998d50b1a67", size = 178592, upload-time = "2024-09-04T20:43:32.108Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/de/cc/4635c320081c78d6ffc2cab0a76025b691a91204f4aa317d568ff9280a2d/cffi-1.17.1-cp310-cp310-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:edae79245293e15384b51f88b00613ba9f7198016a5948b5dddf4917d4d26382", size = 426024, upload-time = "2024-09-04T20:43:34.186Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b6/7b/3b2b250f3aab91abe5f8a51ada1b717935fdaec53f790ad4100fe2ec64d1/cffi-1.17.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:45398b671ac6d70e67da8e4224a065cec6a93541bb7aebe1b198a61b58c7b702", size = 448188, upload-time = "2024-09-04T20:43:36.286Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d3/48/1b9283ebbf0ec065148d8de05d647a986c5f22586b18120020452fff8f5d/cffi-1.17.1-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:ad9413ccdeda48c5afdae7e4fa2192157e991ff761e7ab8fdd8926f40b160cc3", size = 455571, upload-time = "2024-09-04T20:43:38.586Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/40/87/3b8452525437b40f39ca7ff70276679772ee7e8b394934ff60e63b7b090c/cffi-1.17.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:5da5719280082ac6bd9aa7becb3938dc9f9cbd57fac7d2871717b1feb0902ab6", size = 436687, upload-time = "2024-09-04T20:43:40.084Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8d/fb/4da72871d177d63649ac449aec2e8a29efe0274035880c7af59101ca2232/cffi-1.17.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2bb1a08b8008b281856e5971307cc386a8e9c5b625ac297e853d36da6efe9c17", size = 446211, upload-time = "2024-09-04T20:43:41.526Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ab/a0/62f00bcb411332106c02b663b26f3545a9ef136f80d5df746c05878f8c4b/cffi-1.17.1-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:045d61c734659cc045141be4bae381a41d89b741f795af1dd018bfb532fd0df8", size = 461325, upload-time = "2024-09-04T20:43:43.117Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/36/83/76127035ed2e7e27b0787604d99da630ac3123bfb02d8e80c633f218a11d/cffi-1.17.1-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:6883e737d7d9e4899a8a695e00ec36bd4e5e4f18fabe0aca0efe0a4b44cdb13e", size = 438784, upload-time = "2024-09-04T20:43:45.256Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/21/81/a6cd025db2f08ac88b901b745c163d884641909641f9b826e8cb87645942/cffi-1.17.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:6b8b4a92e1c65048ff98cfe1f735ef8f1ceb72e3d5f0c25fdb12087a23da22be", size = 461564, upload-time = "2024-09-04T20:43:46.779Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f8/fe/4d41c2f200c4a457933dbd98d3cf4e911870877bd94d9656cc0fcb390681/cffi-1.17.1-cp310-cp310-win32.whl", hash = "sha256:c9c3d058ebabb74db66e431095118094d06abf53284d9c81f27300d0e0d8bc7c", size = 171804, upload-time = "2024-09-04T20:43:48.186Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d1/b6/0b0f5ab93b0df4acc49cae758c81fe4e5ef26c3ae2e10cc69249dfd8b3ab/cffi-1.17.1-cp310-cp310-win_amd64.whl", hash = "sha256:0f048dcf80db46f0098ccac01132761580d28e28bc0f78ae0d58048063317e15", size = 181299, upload-time = "2024-09-04T20:43:49.812Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6b/f4/927e3a8899e52a27fa57a48607ff7dc91a9ebe97399b357b85a0c7892e00/cffi-1.17.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:a45e3c6913c5b87b3ff120dcdc03f6131fa0065027d0ed7ee6190736a74cd401", size = 182264, upload-time = "2024-09-04T20:43:51.124Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6c/f5/6c3a8efe5f503175aaddcbea6ad0d2c96dad6f5abb205750d1b3df44ef29/cffi-1.17.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:30c5e0cb5ae493c04c8b42916e52ca38079f1b235c2f8ae5f4527b963c401caf", size = 178651, upload-time = "2024-09-04T20:43:52.872Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/94/dd/a3f0118e688d1b1a57553da23b16bdade96d2f9bcda4d32e7d2838047ff7/cffi-1.17.1-cp311-cp311-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f75c7ab1f9e4aca5414ed4d8e5c0e303a34f4421f8a0d47a4d019ceff0ab6af4", size = 445259, upload-time = "2024-09-04T20:43:56.123Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2e/ea/70ce63780f096e16ce8588efe039d3c4f91deb1dc01e9c73a287939c79a6/cffi-1.17.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a1ed2dd2972641495a3ec98445e09766f077aee98a1c896dcb4ad0d303628e41", size = 469200, upload-time = "2024-09-04T20:43:57.891Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1c/a0/a4fa9f4f781bda074c3ddd57a572b060fa0df7655d2a4247bbe277200146/cffi-1.17.1-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:46bf43160c1a35f7ec506d254e5c890f3c03648a4dbac12d624e4490a7046cd1", size = 477235, upload-time = "2024-09-04T20:44:00.18Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/62/12/ce8710b5b8affbcdd5c6e367217c242524ad17a02fe5beec3ee339f69f85/cffi-1.17.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a24ed04c8ffd54b0729c07cee15a81d964e6fee0e3d4d342a27b020d22959dc6", size = 459721, upload-time = "2024-09-04T20:44:01.585Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ff/6b/d45873c5e0242196f042d555526f92aa9e0c32355a1be1ff8c27f077fd37/cffi-1.17.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:610faea79c43e44c71e1ec53a554553fa22321b65fae24889706c0a84d4ad86d", size = 467242, upload-time = "2024-09-04T20:44:03.467Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1a/52/d9a0e523a572fbccf2955f5abe883cfa8bcc570d7faeee06336fbd50c9fc/cffi-1.17.1-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:a9b15d491f3ad5d692e11f6b71f7857e7835eb677955c00cc0aefcd0669adaf6", size = 477999, upload-time = "2024-09-04T20:44:05.023Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/44/74/f2a2460684a1a2d00ca799ad880d54652841a780c4c97b87754f660c7603/cffi-1.17.1-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:de2ea4b5833625383e464549fec1bc395c1bdeeb5f25c4a3a82b5a8c756ec22f", size = 454242, upload-time = "2024-09-04T20:44:06.444Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f8/4a/34599cac7dfcd888ff54e801afe06a19c17787dfd94495ab0c8d35fe99fb/cffi-1.17.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:fc48c783f9c87e60831201f2cce7f3b2e4846bf4d8728eabe54d60700b318a0b", size = 478604, upload-time = "2024-09-04T20:44:08.206Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/34/33/e1b8a1ba29025adbdcda5fb3a36f94c03d771c1b7b12f726ff7fef2ebe36/cffi-1.17.1-cp311-cp311-win32.whl", hash = "sha256:85a950a4ac9c359340d5963966e3e0a94a676bd6245a4b55bc43949eee26a655", size = 171727, upload-time = "2024-09-04T20:44:09.481Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3d/97/50228be003bb2802627d28ec0627837ac0bf35c90cf769812056f235b2d1/cffi-1.17.1-cp311-cp311-win_amd64.whl", hash = "sha256:caaf0640ef5f5517f49bc275eca1406b0ffa6aa184892812030f04c2abf589a0", size = 181400, upload-time = "2024-09-04T20:44:10.873Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5a/84/e94227139ee5fb4d600a7a4927f322e1d4aea6fdc50bd3fca8493caba23f/cffi-1.17.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:805b4371bf7197c329fcb3ead37e710d1bca9da5d583f5073b799d5c5bd1eee4", size = 183178, upload-time = "2024-09-04T20:44:12.232Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/da/ee/fb72c2b48656111c4ef27f0f91da355e130a923473bf5ee75c5643d00cca/cffi-1.17.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:733e99bc2df47476e3848417c5a4540522f234dfd4ef3ab7fafdf555b082ec0c", size = 178840, upload-time = "2024-09-04T20:44:13.739Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/cc/b6/db007700f67d151abadf508cbfd6a1884f57eab90b1bb985c4c8c02b0f28/cffi-1.17.1-cp312-cp312-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1257bdabf294dceb59f5e70c64a3e2f462c30c7ad68092d01bbbfb1c16b1ba36", size = 454803, upload-time = "2024-09-04T20:44:15.231Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1a/df/f8d151540d8c200eb1c6fba8cd0dfd40904f1b0682ea705c36e6c2e97ab3/cffi-1.17.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:da95af8214998d77a98cc14e3a3bd00aa191526343078b530ceb0bd710fb48a5", size = 478850, upload-time = "2024-09-04T20:44:17.188Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/28/c0/b31116332a547fd2677ae5b78a2ef662dfc8023d67f41b2a83f7c2aa78b1/cffi-1.17.1-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:d63afe322132c194cf832bfec0dc69a99fb9bb6bbd550f161a49e9e855cc78ff", size = 485729, upload-time = "2024-09-04T20:44:18.688Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/91/2b/9a1ddfa5c7f13cab007a2c9cc295b70fbbda7cb10a286aa6810338e60ea1/cffi-1.17.1-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:f79fc4fc25f1c8698ff97788206bb3c2598949bfe0fef03d299eb1b5356ada99", size = 471256, upload-time = "2024-09-04T20:44:20.248Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b2/d5/da47df7004cb17e4955df6a43d14b3b4ae77737dff8bf7f8f333196717bf/cffi-1.17.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b62ce867176a75d03a665bad002af8e6d54644fad99a3c70905c543130e39d93", size = 479424, upload-time = "2024-09-04T20:44:21.673Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0b/ac/2a28bcf513e93a219c8a4e8e125534f4f6db03e3179ba1c45e949b76212c/cffi-1.17.1-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:386c8bf53c502fff58903061338ce4f4950cbdcb23e2902d86c0f722b786bbe3", size = 484568, upload-time = "2024-09-04T20:44:23.245Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d4/38/ca8a4f639065f14ae0f1d9751e70447a261f1a30fa7547a828ae08142465/cffi-1.17.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:4ceb10419a9adf4460ea14cfd6bc43d08701f0835e979bf821052f1805850fe8", size = 488736, upload-time = "2024-09-04T20:44:24.757Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/86/c5/28b2d6f799ec0bdecf44dced2ec5ed43e0eb63097b0f58c293583b406582/cffi-1.17.1-cp312-cp312-win32.whl", hash = "sha256:a08d7e755f8ed21095a310a693525137cfe756ce62d066e53f502a83dc550f65", size = 172448, upload-time = "2024-09-04T20:44:26.208Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/50/b9/db34c4755a7bd1cb2d1603ac3863f22bcecbd1ba29e5ee841a4bc510b294/cffi-1.17.1-cp312-cp312-win_amd64.whl", hash = "sha256:51392eae71afec0d0c8fb1a53b204dbb3bcabcb3c9b807eedf3e1e6ccf2de903", size = 181976, upload-time = "2024-09-04T20:44:27.578Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8d/f8/dd6c246b148639254dad4d6803eb6a54e8c85c6e11ec9df2cffa87571dbe/cffi-1.17.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:f3a2b4222ce6b60e2e8b337bb9596923045681d71e5a082783484d845390938e", size = 182989, upload-time = "2024-09-04T20:44:28.956Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8b/f1/672d303ddf17c24fc83afd712316fda78dc6fce1cd53011b839483e1ecc8/cffi-1.17.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:0984a4925a435b1da406122d4d7968dd861c1385afe3b45ba82b750f229811e2", size = 178802, upload-time = "2024-09-04T20:44:30.289Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0e/2d/eab2e858a91fdff70533cab61dcff4a1f55ec60425832ddfdc9cd36bc8af/cffi-1.17.1-cp313-cp313-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d01b12eeeb4427d3110de311e1774046ad344f5b1a7403101878976ecd7a10f3", size = 454792, upload-time = "2024-09-04T20:44:32.01Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/75/b2/fbaec7c4455c604e29388d55599b99ebcc250a60050610fadde58932b7ee/cffi-1.17.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:706510fe141c86a69c8ddc029c7910003a17353970cff3b904ff0686a5927683", size = 478893, upload-time = "2024-09-04T20:44:33.606Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4f/b7/6e4a2162178bf1935c336d4da8a9352cccab4d3a5d7914065490f08c0690/cffi-1.17.1-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:de55b766c7aa2e2a3092c51e0483d700341182f08e67c63630d5b6f200bb28e5", size = 485810, upload-time = "2024-09-04T20:44:35.191Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c7/8a/1d0e4a9c26e54746dc08c2c6c037889124d4f59dffd853a659fa545f1b40/cffi-1.17.1-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:c59d6e989d07460165cc5ad3c61f9fd8f1b4796eacbd81cee78957842b834af4", size = 471200, upload-time = "2024-09-04T20:44:36.743Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/26/9f/1aab65a6c0db35f43c4d1b4f580e8df53914310afc10ae0397d29d697af4/cffi-1.17.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dd398dbc6773384a17fe0d3e7eeb8d1a21c2200473ee6806bb5e6a8e62bb73dd", size = 479447, upload-time = "2024-09-04T20:44:38.492Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5f/e4/fb8b3dd8dc0e98edf1135ff067ae070bb32ef9d509d6cb0f538cd6f7483f/cffi-1.17.1-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:3edc8d958eb099c634dace3c7e16560ae474aa3803a5df240542b305d14e14ed", size = 484358, upload-time = "2024-09-04T20:44:40.046Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f1/47/d7145bf2dc04684935d57d67dff9d6d795b2ba2796806bb109864be3a151/cffi-1.17.1-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:72e72408cad3d5419375fc87d289076ee319835bdfa2caad331e377589aebba9", size = 488469, upload-time = "2024-09-04T20:44:41.616Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/bf/ee/f94057fa6426481d663b88637a9a10e859e492c73d0384514a17d78ee205/cffi-1.17.1-cp313-cp313-win32.whl", hash = "sha256:e03eab0a8677fa80d646b5ddece1cbeaf556c313dcfac435ba11f107ba117b5d", size = 172475, upload-time = "2024-09-04T20:44:43.733Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7c/fc/6a8cb64e5f0324877d503c854da15d76c1e50eb722e320b15345c4d0c6de/cffi-1.17.1-cp313-cp313-win_amd64.whl", hash = "sha256:f6a16c31041f09ead72d69f583767292f750d24913dadacf5756b966aacb3f1a", size = 182009, upload-time = "2024-09-04T20:44:45.309Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/48/08/15bf6b43ae9bd06f6b00ad8a91f5a8fe1069d4c9fab550a866755402724e/cffi-1.17.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:636062ea65bd0195bc012fea9321aca499c0504409f413dc88af450b57ffd03b", size = 182457, upload-time = "2024-09-04T20:44:47.892Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c2/5b/f1523dd545f92f7df468e5f653ffa4df30ac222f3c884e51e139878f1cb5/cffi-1.17.1-cp38-cp38-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c7eac2ef9b63c79431bc4b25f1cd649d7f061a28808cbc6c47b534bd789ef964", size = 425932, upload-time = "2024-09-04T20:44:49.491Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/53/93/7e547ab4105969cc8c93b38a667b82a835dd2cc78f3a7dad6130cfd41e1d/cffi-1.17.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e221cf152cff04059d011ee126477f0d9588303eb57e88923578ace7baad17f9", size = 448585, upload-time = "2024-09-04T20:44:51.671Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/56/c4/a308f2c332006206bb511de219efeff090e9d63529ba0a77aae72e82248b/cffi-1.17.1-cp38-cp38-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:31000ec67d4221a71bd3f67df918b1f88f676f1c3b535a7eb473255fdc0b83fc", size = 456268, upload-time = "2024-09-04T20:44:53.51Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ca/5b/b63681518265f2f4060d2b60755c1c77ec89e5e045fc3773b72735ddaad5/cffi-1.17.1-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:6f17be4345073b0a7b8ea599688f692ac3ef23ce28e5df79c04de519dbc4912c", size = 436592, upload-time = "2024-09-04T20:44:55.085Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/bb/19/b51af9f4a4faa4a8ac5a0e5d5c2522dcd9703d07fac69da34a36c4d960d3/cffi-1.17.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0e2b1fac190ae3ebfe37b979cc1ce69c81f4e4fe5746bb401dca63a9062cdaf1", size = 446512, upload-time = "2024-09-04T20:44:57.135Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e2/63/2bed8323890cb613bbecda807688a31ed11a7fe7afe31f8faaae0206a9a3/cffi-1.17.1-cp38-cp38-win32.whl", hash = "sha256:7596d6620d3fa590f677e9ee430df2958d2d6d6de2feeae5b20e82c00b76fbf8", size = 171576, upload-time = "2024-09-04T20:44:58.535Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2f/70/80c33b044ebc79527447fd4fbc5455d514c3bb840dede4455de97da39b4d/cffi-1.17.1-cp38-cp38-win_amd64.whl", hash = "sha256:78122be759c3f8a014ce010908ae03364d00a1f81ab5c7f4a7a5120607ea56e1", size = 181229, upload-time = "2024-09-04T20:44:59.963Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b9/ea/8bb50596b8ffbc49ddd7a1ad305035daa770202a6b782fc164647c2673ad/cffi-1.17.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:b2ab587605f4ba0bf81dc0cb08a41bd1c0a5906bd59243d56bad7668a6fc6c16", size = 182220, upload-time = "2024-09-04T20:45:01.577Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ae/11/e77c8cd24f58285a82c23af484cf5b124a376b32644e445960d1a4654c3a/cffi-1.17.1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:28b16024becceed8c6dfbc75629e27788d8a3f9030691a1dbf9821a128b22c36", size = 178605, upload-time = "2024-09-04T20:45:03.837Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ed/65/25a8dc32c53bf5b7b6c2686b42ae2ad58743f7ff644844af7cdb29b49361/cffi-1.17.1-cp39-cp39-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1d599671f396c4723d016dbddb72fe8e0397082b0a77a4fab8028923bec050e8", size = 424910, upload-time = "2024-09-04T20:45:05.315Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/42/7a/9d086fab7c66bd7c4d0f27c57a1b6b068ced810afc498cc8c49e0088661c/cffi-1.17.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ca74b8dbe6e8e8263c0ffd60277de77dcee6c837a3d0881d8c1ead7268c9e576", size = 447200, upload-time = "2024-09-04T20:45:06.903Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/da/63/1785ced118ce92a993b0ec9e0d0ac8dc3e5dbfbcaa81135be56c69cabbb6/cffi-1.17.1-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:f7f5baafcc48261359e14bcd6d9bff6d4b28d9103847c9e136694cb0501aef87", size = 454565, upload-time = "2024-09-04T20:45:08.975Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/74/06/90b8a44abf3556599cdec107f7290277ae8901a58f75e6fe8f970cd72418/cffi-1.17.1-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:98e3969bcff97cae1b2def8ba499ea3d6f31ddfdb7635374834cf89a1a08ecf0", size = 435635, upload-time = "2024-09-04T20:45:10.64Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/bd/62/a1f468e5708a70b1d86ead5bab5520861d9c7eacce4a885ded9faa7729c3/cffi-1.17.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:cdf5ce3acdfd1661132f2a9c19cac174758dc2352bfe37d98aa7512c6b7178b3", size = 445218, upload-time = "2024-09-04T20:45:12.366Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5b/95/b34462f3ccb09c2594aa782d90a90b045de4ff1f70148ee79c69d37a0a5a/cffi-1.17.1-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:9755e4345d1ec879e3849e62222a18c7174d65a6a92d5b346b1863912168b595", size = 460486, upload-time = "2024-09-04T20:45:13.935Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fc/fc/a1e4bebd8d680febd29cf6c8a40067182b64f00c7d105f8f26b5bc54317b/cffi-1.17.1-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:f1e22e8c4419538cb197e4dd60acc919d7696e5ef98ee4da4e01d3f8cfa4cc5a", size = 437911, upload-time = "2024-09-04T20:45:15.696Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e6/c3/21cab7a6154b6a5ea330ae80de386e7665254835b9e98ecc1340b3a7de9a/cffi-1.17.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:c03e868a0b3bc35839ba98e74211ed2b05d2119be4e8a0f224fba9384f1fe02e", size = 460632, upload-time = "2024-09-04T20:45:17.284Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/cb/b5/fd9f8b5a84010ca169ee49f4e4ad6f8c05f4e3545b72ee041dbbcb159882/cffi-1.17.1-cp39-cp39-win32.whl", hash = "sha256:e31ae45bc2e29f6b2abd0de1cc3b9d5205aa847cafaecb8af1476a609a2f6eb7", size = 171820, upload-time = "2024-09-04T20:45:18.762Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8c/52/b08750ce0bce45c143e1b5d7357ee8c55341b52bdef4b0f081af1eb248c2/cffi-1.17.1-cp39-cp39-win_amd64.whl", hash = "sha256:d016c76bdd850f3c626af19b0542c9677ba156e4ee4fccfdd7848803533ef662", size = 181290, upload-time = "2024-09-04T20:45:20.226Z" }, +] + +[[package]] +name = "charset-normalizer" +version = "3.4.3" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/83/2d/5fd176ceb9b2fc619e63405525573493ca23441330fcdaee6bef9460e924/charset_normalizer-3.4.3.tar.gz", hash = "sha256:6fce4b8500244f6fcb71465d4a4930d132ba9ab8e71a7859e6a5d59851068d14", size = 122371, upload-time = "2025-08-09T07:57:28.46Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d6/98/f3b8013223728a99b908c9344da3aa04ee6e3fa235f19409033eda92fb78/charset_normalizer-3.4.3-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:fb7f67a1bfa6e40b438170ebdc8158b78dc465a5a67b6dde178a46987b244a72", size = 207695, upload-time = "2025-08-09T07:55:36.452Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/21/40/5188be1e3118c82dcb7c2a5ba101b783822cfb413a0268ed3be0468532de/charset_normalizer-3.4.3-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cc9370a2da1ac13f0153780040f465839e6cccb4a1e44810124b4e22483c93fe", size = 147153, upload-time = "2025-08-09T07:55:38.467Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/37/60/5d0d74bc1e1380f0b72c327948d9c2aca14b46a9efd87604e724260f384c/charset_normalizer-3.4.3-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:07a0eae9e2787b586e129fdcbe1af6997f8d0e5abaa0bc98c0e20e124d67e601", size = 160428, upload-time = "2025-08-09T07:55:40.072Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/85/9a/d891f63722d9158688de58d050c59dc3da560ea7f04f4c53e769de5140f5/charset_normalizer-3.4.3-cp310-cp310-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:74d77e25adda8581ffc1c720f1c81ca082921329452eba58b16233ab1842141c", size = 157627, upload-time = "2025-08-09T07:55:41.706Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/65/1a/7425c952944a6521a9cfa7e675343f83fd82085b8af2b1373a2409c683dc/charset_normalizer-3.4.3-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d0e909868420b7049dafd3a31d45125b31143eec59235311fc4c57ea26a4acd2", size = 152388, upload-time = "2025-08-09T07:55:43.262Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f0/c9/a2c9c2a355a8594ce2446085e2ec97fd44d323c684ff32042e2a6b718e1d/charset_normalizer-3.4.3-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:c6f162aabe9a91a309510d74eeb6507fab5fff92337a15acbe77753d88d9dcf0", size = 150077, upload-time = "2025-08-09T07:55:44.903Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3b/38/20a1f44e4851aa1c9105d6e7110c9d020e093dfa5836d712a5f074a12bf7/charset_normalizer-3.4.3-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:4ca4c094de7771a98d7fbd67d9e5dbf1eb73efa4f744a730437d8a3a5cf994f0", size = 161631, upload-time = "2025-08-09T07:55:46.346Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a4/fa/384d2c0f57edad03d7bec3ebefb462090d8905b4ff5a2d2525f3bb711fac/charset_normalizer-3.4.3-cp310-cp310-musllinux_1_2_s390x.whl", hash = "sha256:02425242e96bcf29a49711b0ca9f37e451da7c70562bc10e8ed992a5a7a25cc0", size = 159210, upload-time = "2025-08-09T07:55:47.539Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/33/9e/eca49d35867ca2db336b6ca27617deed4653b97ebf45dfc21311ce473c37/charset_normalizer-3.4.3-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:78deba4d8f9590fe4dae384aeff04082510a709957e968753ff3c48399f6f92a", size = 153739, upload-time = "2025-08-09T07:55:48.744Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2a/91/26c3036e62dfe8de8061182d33be5025e2424002125c9500faff74a6735e/charset_normalizer-3.4.3-cp310-cp310-win32.whl", hash = "sha256:d79c198e27580c8e958906f803e63cddb77653731be08851c7df0b1a14a8fc0f", size = 99825, upload-time = "2025-08-09T07:55:50.305Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e2/c6/f05db471f81af1fa01839d44ae2a8bfeec8d2a8b4590f16c4e7393afd323/charset_normalizer-3.4.3-cp310-cp310-win_amd64.whl", hash = "sha256:c6e490913a46fa054e03699c70019ab869e990270597018cef1d8562132c2669", size = 107452, upload-time = "2025-08-09T07:55:51.461Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7f/b5/991245018615474a60965a7c9cd2b4efbaabd16d582a5547c47ee1c7730b/charset_normalizer-3.4.3-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:b256ee2e749283ef3ddcff51a675ff43798d92d746d1a6e4631bf8c707d22d0b", size = 204483, upload-time = "2025-08-09T07:55:53.12Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c7/2a/ae245c41c06299ec18262825c1569c5d3298fc920e4ddf56ab011b417efd/charset_normalizer-3.4.3-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:13faeacfe61784e2559e690fc53fa4c5ae97c6fcedb8eb6fb8d0a15b475d2c64", size = 145520, upload-time = "2025-08-09T07:55:54.712Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3a/a4/b3b6c76e7a635748c4421d2b92c7b8f90a432f98bda5082049af37ffc8e3/charset_normalizer-3.4.3-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:00237675befef519d9af72169d8604a067d92755e84fe76492fef5441db05b91", size = 158876, upload-time = "2025-08-09T07:55:56.024Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e2/e6/63bb0e10f90a8243c5def74b5b105b3bbbfb3e7bb753915fe333fb0c11ea/charset_normalizer-3.4.3-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:585f3b2a80fbd26b048a0be90c5aae8f06605d3c92615911c3a2b03a8a3b796f", size = 156083, upload-time = "2025-08-09T07:55:57.582Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/87/df/b7737ff046c974b183ea9aa111b74185ac8c3a326c6262d413bd5a1b8c69/charset_normalizer-3.4.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0e78314bdc32fa80696f72fa16dc61168fda4d6a0c014e0380f9d02f0e5d8a07", size = 150295, upload-time = "2025-08-09T07:55:59.147Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/61/f1/190d9977e0084d3f1dc169acd060d479bbbc71b90bf3e7bf7b9927dec3eb/charset_normalizer-3.4.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:96b2b3d1a83ad55310de8c7b4a2d04d9277d5591f40761274856635acc5fcb30", size = 148379, upload-time = "2025-08-09T07:56:00.364Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4c/92/27dbe365d34c68cfe0ca76f1edd70e8705d82b378cb54ebbaeabc2e3029d/charset_normalizer-3.4.3-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:939578d9d8fd4299220161fdd76e86c6a251987476f5243e8864a7844476ba14", size = 160018, upload-time = "2025-08-09T07:56:01.678Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/99/04/baae2a1ea1893a01635d475b9261c889a18fd48393634b6270827869fa34/charset_normalizer-3.4.3-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:fd10de089bcdcd1be95a2f73dbe6254798ec1bda9f450d5828c96f93e2536b9c", size = 157430, upload-time = "2025-08-09T07:56:02.87Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2f/36/77da9c6a328c54d17b960c89eccacfab8271fdaaa228305330915b88afa9/charset_normalizer-3.4.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:1e8ac75d72fa3775e0b7cb7e4629cec13b7514d928d15ef8ea06bca03ef01cae", size = 151600, upload-time = "2025-08-09T07:56:04.089Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/64/d4/9eb4ff2c167edbbf08cdd28e19078bf195762e9bd63371689cab5ecd3d0d/charset_normalizer-3.4.3-cp311-cp311-win32.whl", hash = "sha256:6cf8fd4c04756b6b60146d98cd8a77d0cdae0e1ca20329da2ac85eed779b6849", size = 99616, upload-time = "2025-08-09T07:56:05.658Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f4/9c/996a4a028222e7761a96634d1820de8a744ff4327a00ada9c8942033089b/charset_normalizer-3.4.3-cp311-cp311-win_amd64.whl", hash = "sha256:31a9a6f775f9bcd865d88ee350f0ffb0e25936a7f930ca98995c05abf1faf21c", size = 107108, upload-time = "2025-08-09T07:56:07.176Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e9/5e/14c94999e418d9b87682734589404a25854d5f5d0408df68bc15b6ff54bb/charset_normalizer-3.4.3-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:e28e334d3ff134e88989d90ba04b47d84382a828c061d0d1027b1b12a62b39b1", size = 205655, upload-time = "2025-08-09T07:56:08.475Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7d/a8/c6ec5d389672521f644505a257f50544c074cf5fc292d5390331cd6fc9c3/charset_normalizer-3.4.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0cacf8f7297b0c4fcb74227692ca46b4a5852f8f4f24b3c766dd94a1075c4884", size = 146223, upload-time = "2025-08-09T07:56:09.708Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fc/eb/a2ffb08547f4e1e5415fb69eb7db25932c52a52bed371429648db4d84fb1/charset_normalizer-3.4.3-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c6fd51128a41297f5409deab284fecbe5305ebd7e5a1f959bee1c054622b7018", size = 159366, upload-time = "2025-08-09T07:56:11.326Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/82/10/0fd19f20c624b278dddaf83b8464dcddc2456cb4b02bb902a6da126b87a1/charset_normalizer-3.4.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:3cfb2aad70f2c6debfbcb717f23b7eb55febc0bb23dcffc0f076009da10c6392", size = 157104, upload-time = "2025-08-09T07:56:13.014Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/16/ab/0233c3231af734f5dfcf0844aa9582d5a1466c985bbed6cedab85af9bfe3/charset_normalizer-3.4.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1606f4a55c0fd363d754049cdf400175ee96c992b1f8018b993941f221221c5f", size = 151830, upload-time = "2025-08-09T07:56:14.428Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ae/02/e29e22b4e02839a0e4a06557b1999d0a47db3567e82989b5bb21f3fbbd9f/charset_normalizer-3.4.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:027b776c26d38b7f15b26a5da1044f376455fb3766df8fc38563b4efbc515154", size = 148854, upload-time = "2025-08-09T07:56:16.051Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/05/6b/e2539a0a4be302b481e8cafb5af8792da8093b486885a1ae4d15d452bcec/charset_normalizer-3.4.3-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:42e5088973e56e31e4fa58eb6bd709e42fc03799c11c42929592889a2e54c491", size = 160670, upload-time = "2025-08-09T07:56:17.314Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/31/e7/883ee5676a2ef217a40ce0bffcc3d0dfbf9e64cbcfbdf822c52981c3304b/charset_normalizer-3.4.3-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:cc34f233c9e71701040d772aa7490318673aa7164a0efe3172b2981218c26d93", size = 158501, upload-time = "2025-08-09T07:56:18.641Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c1/35/6525b21aa0db614cf8b5792d232021dca3df7f90a1944db934efa5d20bb1/charset_normalizer-3.4.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:320e8e66157cc4e247d9ddca8e21f427efc7a04bbd0ac8a9faf56583fa543f9f", size = 153173, upload-time = "2025-08-09T07:56:20.289Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/50/ee/f4704bad8201de513fdc8aac1cabc87e38c5818c93857140e06e772b5892/charset_normalizer-3.4.3-cp312-cp312-win32.whl", hash = "sha256:fb6fecfd65564f208cbf0fba07f107fb661bcd1a7c389edbced3f7a493f70e37", size = 99822, upload-time = "2025-08-09T07:56:21.551Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/39/f5/3b3836ca6064d0992c58c7561c6b6eee1b3892e9665d650c803bd5614522/charset_normalizer-3.4.3-cp312-cp312-win_amd64.whl", hash = "sha256:86df271bf921c2ee3818f0522e9a5b8092ca2ad8b065ece5d7d9d0e9f4849bcc", size = 107543, upload-time = "2025-08-09T07:56:23.115Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/65/ca/2135ac97709b400c7654b4b764daf5c5567c2da45a30cdd20f9eefe2d658/charset_normalizer-3.4.3-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:14c2a87c65b351109f6abfc424cab3927b3bdece6f706e4d12faaf3d52ee5efe", size = 205326, upload-time = "2025-08-09T07:56:24.721Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/71/11/98a04c3c97dd34e49c7d247083af03645ca3730809a5509443f3c37f7c99/charset_normalizer-3.4.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:41d1fc408ff5fdfb910200ec0e74abc40387bccb3252f3f27c0676731df2b2c8", size = 146008, upload-time = "2025-08-09T07:56:26.004Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/60/f5/4659a4cb3c4ec146bec80c32d8bb16033752574c20b1252ee842a95d1a1e/charset_normalizer-3.4.3-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:1bb60174149316da1c35fa5233681f7c0f9f514509b8e399ab70fea5f17e45c9", size = 159196, upload-time = "2025-08-09T07:56:27.25Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/86/9e/f552f7a00611f168b9a5865a1414179b2c6de8235a4fa40189f6f79a1753/charset_normalizer-3.4.3-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:30d006f98569de3459c2fc1f2acde170b7b2bd265dc1943e87e1a4efe1b67c31", size = 156819, upload-time = "2025-08-09T07:56:28.515Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7e/95/42aa2156235cbc8fa61208aded06ef46111c4d3f0de233107b3f38631803/charset_normalizer-3.4.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:416175faf02e4b0810f1f38bcb54682878a4af94059a1cd63b8747244420801f", size = 151350, upload-time = "2025-08-09T07:56:29.716Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c2/a9/3865b02c56f300a6f94fc631ef54f0a8a29da74fb45a773dfd3dcd380af7/charset_normalizer-3.4.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:6aab0f181c486f973bc7262a97f5aca3ee7e1437011ef0c2ec04b5a11d16c927", size = 148644, upload-time = "2025-08-09T07:56:30.984Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/77/d9/cbcf1a2a5c7d7856f11e7ac2d782aec12bdfea60d104e60e0aa1c97849dc/charset_normalizer-3.4.3-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:fdabf8315679312cfa71302f9bd509ded4f2f263fb5b765cf1433b39106c3cc9", size = 160468, upload-time = "2025-08-09T07:56:32.252Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f6/42/6f45efee8697b89fda4d50580f292b8f7f9306cb2971d4b53f8914e4d890/charset_normalizer-3.4.3-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:bd28b817ea8c70215401f657edef3a8aa83c29d447fb0b622c35403780ba11d5", size = 158187, upload-time = "2025-08-09T07:56:33.481Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/70/99/f1c3bdcfaa9c45b3ce96f70b14f070411366fa19549c1d4832c935d8e2c3/charset_normalizer-3.4.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:18343b2d246dc6761a249ba1fb13f9ee9a2bcd95decc767319506056ea4ad4dc", size = 152699, upload-time = "2025-08-09T07:56:34.739Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a3/ad/b0081f2f99a4b194bcbb1934ef3b12aa4d9702ced80a37026b7607c72e58/charset_normalizer-3.4.3-cp313-cp313-win32.whl", hash = "sha256:6fb70de56f1859a3f71261cbe41005f56a7842cc348d3aeb26237560bfa5e0ce", size = 99580, upload-time = "2025-08-09T07:56:35.981Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/9a/8f/ae790790c7b64f925e5c953b924aaa42a243fb778fed9e41f147b2a5715a/charset_normalizer-3.4.3-cp313-cp313-win_amd64.whl", hash = "sha256:cf1ebb7d78e1ad8ec2a8c4732c7be2e736f6e5123a4146c5b89c9d1f585f8cef", size = 107366, upload-time = "2025-08-09T07:56:37.339Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8e/91/b5a06ad970ddc7a0e513112d40113e834638f4ca1120eb727a249fb2715e/charset_normalizer-3.4.3-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:3cd35b7e8aedeb9e34c41385fda4f73ba609e561faedfae0a9e75e44ac558a15", size = 204342, upload-time = "2025-08-09T07:56:38.687Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ce/ec/1edc30a377f0a02689342f214455c3f6c2fbedd896a1d2f856c002fc3062/charset_normalizer-3.4.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b89bc04de1d83006373429975f8ef9e7932534b8cc9ca582e4db7d20d91816db", size = 145995, upload-time = "2025-08-09T07:56:40.048Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/17/e5/5e67ab85e6d22b04641acb5399c8684f4d37caf7558a53859f0283a650e9/charset_normalizer-3.4.3-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:2001a39612b241dae17b4687898843f254f8748b796a2e16f1051a17078d991d", size = 158640, upload-time = "2025-08-09T07:56:41.311Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f1/e5/38421987f6c697ee3722981289d554957c4be652f963d71c5e46a262e135/charset_normalizer-3.4.3-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:8dcfc373f888e4fb39a7bc57e93e3b845e7f462dacc008d9749568b1c4ece096", size = 156636, upload-time = "2025-08-09T07:56:43.195Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a0/e4/5a075de8daa3ec0745a9a3b54467e0c2967daaaf2cec04c845f73493e9a1/charset_normalizer-3.4.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:18b97b8404387b96cdbd30ad660f6407799126d26a39ca65729162fd810a99aa", size = 150939, upload-time = "2025-08-09T07:56:44.819Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/02/f7/3611b32318b30974131db62b4043f335861d4d9b49adc6d57c1149cc49d4/charset_normalizer-3.4.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ccf600859c183d70eb47e05a44cd80a4ce77394d1ac0f79dbd2dd90a69a3a049", size = 148580, upload-time = "2025-08-09T07:56:46.684Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7e/61/19b36f4bd67f2793ab6a99b979b4e4f3d8fc754cbdffb805335df4337126/charset_normalizer-3.4.3-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:53cd68b185d98dde4ad8990e56a58dea83a4162161b1ea9272e5c9182ce415e0", size = 159870, upload-time = "2025-08-09T07:56:47.941Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/06/57/84722eefdd338c04cf3030ada66889298eaedf3e7a30a624201e0cbe424a/charset_normalizer-3.4.3-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:30a96e1e1f865f78b030d65241c1ee850cdf422d869e9028e2fc1d5e4db73b92", size = 157797, upload-time = "2025-08-09T07:56:49.756Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/72/2a/aff5dd112b2f14bcc3462c312dce5445806bfc8ab3a7328555da95330e4b/charset_normalizer-3.4.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:d716a916938e03231e86e43782ca7878fb602a125a91e7acb8b5112e2e96ac16", size = 152224, upload-time = "2025-08-09T07:56:51.369Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b7/8c/9839225320046ed279c6e839d51f028342eb77c91c89b8ef2549f951f3ec/charset_normalizer-3.4.3-cp314-cp314-win32.whl", hash = "sha256:c6dbd0ccdda3a2ba7c2ecd9d77b37f3b5831687d8dc1b6ca5f56a4880cc7b7ce", size = 100086, upload-time = "2025-08-09T07:56:52.722Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ee/7a/36fbcf646e41f710ce0a563c1c9a343c6edf9be80786edeb15b6f62e17db/charset_normalizer-3.4.3-cp314-cp314-win_amd64.whl", hash = "sha256:73dc19b562516fc9bcf6e5d6e596df0b4eb98d87e4f79f3ae71840e6ed21361c", size = 107400, upload-time = "2025-08-09T07:56:55.172Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/22/82/63a45bfc36f73efe46731a3a71cb84e2112f7e0b049507025ce477f0f052/charset_normalizer-3.4.3-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:0f2be7e0cf7754b9a30eb01f4295cc3d4358a479843b31f328afd210e2c7598c", size = 198805, upload-time = "2025-08-09T07:56:56.496Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0c/52/8b0c6c3e53f7e546a5e49b9edb876f379725914e1130297f3b423c7b71c5/charset_normalizer-3.4.3-cp38-cp38-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c60e092517a73c632ec38e290eba714e9627abe9d301c8c8a12ec32c314a2a4b", size = 142862, upload-time = "2025-08-09T07:56:57.751Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/59/c0/a74f3bd167d311365e7973990243f32c35e7a94e45103125275b9e6c479f/charset_normalizer-3.4.3-cp38-cp38-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:252098c8c7a873e17dd696ed98bbe91dbacd571da4b87df3736768efa7a792e4", size = 155104, upload-time = "2025-08-09T07:56:58.984Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1a/79/ae516e678d6e32df2e7e740a7be51dc80b700e2697cb70054a0f1ac2c955/charset_normalizer-3.4.3-cp38-cp38-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:3653fad4fe3ed447a596ae8638b437f827234f01a8cd801842e43f3d0a6b281b", size = 152598, upload-time = "2025-08-09T07:57:00.201Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/00/bd/ef9c88464b126fa176f4ef4a317ad9b6f4d30b2cffbc43386062367c3e2c/charset_normalizer-3.4.3-cp38-cp38-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8999f965f922ae054125286faf9f11bc6932184b93011d138925a1773830bbe9", size = 147391, upload-time = "2025-08-09T07:57:01.441Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7a/03/cbb6fac9d3e57f7e07ce062712ee80d80a5ab46614684078461917426279/charset_normalizer-3.4.3-cp38-cp38-musllinux_1_2_aarch64.whl", hash = "sha256:d95bfb53c211b57198bb91c46dd5a2d8018b3af446583aab40074bf7988401cb", size = 145037, upload-time = "2025-08-09T07:57:02.638Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/64/d1/f9d141c893ef5d4243bc75c130e95af8fd4bc355beff06e9b1e941daad6e/charset_normalizer-3.4.3-cp38-cp38-musllinux_1_2_ppc64le.whl", hash = "sha256:5b413b0b1bfd94dbf4023ad6945889f374cd24e3f62de58d6bb102c4d9ae534a", size = 156425, upload-time = "2025-08-09T07:57:03.898Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c5/35/9c99739250742375167bc1b1319cd1cec2bf67438a70d84b2e1ec4c9daa3/charset_normalizer-3.4.3-cp38-cp38-musllinux_1_2_s390x.whl", hash = "sha256:b5e3b2d152e74e100a9e9573837aba24aab611d39428ded46f4e4022ea7d1942", size = 153734, upload-time = "2025-08-09T07:57:05.549Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/50/10/c117806094d2c956ba88958dab680574019abc0c02bcf57b32287afca544/charset_normalizer-3.4.3-cp38-cp38-musllinux_1_2_x86_64.whl", hash = "sha256:a2d08ac246bb48479170408d6c19f6385fa743e7157d716e144cad849b2dd94b", size = 148551, upload-time = "2025-08-09T07:57:06.823Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/61/c5/dc3ba772489c453621ffc27e8978a98fe7e41a93e787e5e5bde797f1dddb/charset_normalizer-3.4.3-cp38-cp38-win32.whl", hash = "sha256:ec557499516fc90fd374bf2e32349a2887a876fbf162c160e3c01b6849eaf557", size = 98459, upload-time = "2025-08-09T07:57:08.031Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/05/35/bb59b1cd012d7196fc81c2f5879113971efc226a63812c9cf7f89fe97c40/charset_normalizer-3.4.3-cp38-cp38-win_amd64.whl", hash = "sha256:5d8d01eac18c423815ed4f4a2ec3b439d654e55ee4ad610e153cf02faf67ea40", size = 105887, upload-time = "2025-08-09T07:57:09.401Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c2/ca/9a0983dd5c8e9733565cf3db4df2b0a2e9a82659fd8aa2a868ac6e4a991f/charset_normalizer-3.4.3-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:70bfc5f2c318afece2f5838ea5e4c3febada0be750fcf4775641052bbba14d05", size = 207520, upload-time = "2025-08-09T07:57:11.026Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/39/c6/99271dc37243a4f925b09090493fb96c9333d7992c6187f5cfe5312008d2/charset_normalizer-3.4.3-cp39-cp39-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:23b6b24d74478dc833444cbd927c338349d6ae852ba53a0d02a2de1fce45b96e", size = 147307, upload-time = "2025-08-09T07:57:12.4Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e4/69/132eab043356bba06eb333cc2cc60c6340857d0a2e4ca6dc2b51312886b3/charset_normalizer-3.4.3-cp39-cp39-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:34a7f768e3f985abdb42841e20e17b330ad3aaf4bb7e7aeeb73db2e70f077b99", size = 160448, upload-time = "2025-08-09T07:57:13.712Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/04/9a/914d294daa4809c57667b77470533e65def9c0be1ef8b4c1183a99170e9d/charset_normalizer-3.4.3-cp39-cp39-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:fb731e5deb0c7ef82d698b0f4c5bb724633ee2a489401594c5c88b02e6cb15f7", size = 157758, upload-time = "2025-08-09T07:57:14.979Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b0/a8/6f5bcf1bcf63cb45625f7c5cadca026121ff8a6c8a3256d8d8cd59302663/charset_normalizer-3.4.3-cp39-cp39-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:257f26fed7d7ff59921b78244f3cd93ed2af1800ff048c33f624c87475819dd7", size = 152487, upload-time = "2025-08-09T07:57:16.332Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c4/72/d3d0e9592f4e504f9dea08b8db270821c909558c353dc3b457ed2509f2fb/charset_normalizer-3.4.3-cp39-cp39-musllinux_1_2_aarch64.whl", hash = "sha256:1ef99f0456d3d46a50945c98de1774da86f8e992ab5c77865ea8b8195341fc19", size = 150054, upload-time = "2025-08-09T07:57:17.576Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/20/30/5f64fe3981677fe63fa987b80e6c01042eb5ff653ff7cec1b7bd9268e54e/charset_normalizer-3.4.3-cp39-cp39-musllinux_1_2_ppc64le.whl", hash = "sha256:2c322db9c8c89009a990ef07c3bcc9f011a3269bc06782f916cd3d9eed7c9312", size = 161703, upload-time = "2025-08-09T07:57:20.012Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e1/ef/dd08b2cac9284fd59e70f7d97382c33a3d0a926e45b15fc21b3308324ffd/charset_normalizer-3.4.3-cp39-cp39-musllinux_1_2_s390x.whl", hash = "sha256:511729f456829ef86ac41ca78c63a5cb55240ed23b4b737faca0eb1abb1c41bc", size = 159096, upload-time = "2025-08-09T07:57:21.329Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/45/8c/dcef87cfc2b3f002a6478f38906f9040302c68aebe21468090e39cde1445/charset_normalizer-3.4.3-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:88ab34806dea0671532d3f82d82b85e8fc23d7b2dd12fa837978dad9bb392a34", size = 153852, upload-time = "2025-08-09T07:57:22.608Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/63/86/9cbd533bd37883d467fcd1bd491b3547a3532d0fbb46de2b99feeebf185e/charset_normalizer-3.4.3-cp39-cp39-win32.whl", hash = "sha256:16a8770207946ac75703458e2c743631c79c59c5890c80011d536248f8eaa432", size = 99840, upload-time = "2025-08-09T07:57:23.883Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ce/d6/7e805c8e5c46ff9729c49950acc4ee0aeb55efb8b3a56687658ad10c3216/charset_normalizer-3.4.3-cp39-cp39-win_amd64.whl", hash = "sha256:d22dbedd33326a4a5190dd4fe9e9e693ef12160c77382d9e87919bce54f3d4ca", size = 107438, upload-time = "2025-08-09T07:57:25.287Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8a/1f/f041989e93b001bc4e44bb1669ccdcf54d3f00e628229a85b08d330615c5/charset_normalizer-3.4.3-py3-none-any.whl", hash = "sha256:ce571ab16d890d23b5c278547ba694193a45011ff86a9162a71307ed9f86759a", size = 53175, upload-time = "2025-08-09T07:57:26.864Z" }, +] + +[[package]] +name = "colorama" +version = "0.4.6" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, +] + +[[package]] +name = "decorator" +version = "5.2.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/43/fa/6d96a0978d19e17b68d634497769987b16c8f4cd0a7a05048bec693caa6b/decorator-5.2.1.tar.gz", hash = "sha256:65f266143752f734b0a7cc83c46f4618af75b8c5911b00ccb61d0ac9b6da0360", size = 56711, upload-time = "2025-02-24T04:41:34.073Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4e/8c/f3147f5c4b73e7550fe5f9352eaa956ae838d5c51eb58e7a25b9f3e2643b/decorator-5.2.1-py3-none-any.whl", hash = "sha256:d316bb415a2d9e2d2b3abcc4084c6502fc09240e292cd76a76afc106a1c8e04a", size = 9190, upload-time = "2025-02-24T04:41:32.565Z" }, +] + +[[package]] +name = "einops" +version = "0.8.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e5/81/df4fbe24dff8ba3934af99044188e20a98ed441ad17a274539b74e82e126/einops-0.8.1.tar.gz", hash = "sha256:de5d960a7a761225532e0f1959e5315ebeafc0cd43394732f103ca44b9837e84", size = 54805, upload-time = "2025-02-09T03:17:00.434Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/87/62/9773de14fe6c45c23649e98b83231fffd7b9892b6cf863251dc2afa73643/einops-0.8.1-py3-none-any.whl", hash = "sha256:919387eb55330f5757c6bea9165c5ff5cfe63a642682ea788a6d472576d81737", size = 64359, upload-time = "2025-02-09T03:17:01.998Z" }, +] + +[[package]] +name = "et-xmlfile" +version = "2.0.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d3/38/af70d7ab1ae9d4da450eeec1fa3918940a5fafb9055e934af8d6eb0c2313/et_xmlfile-2.0.0.tar.gz", hash = "sha256:dab3f4764309081ce75662649be815c4c9081e88f0837825f90fd28317d4da54", size = 17234, upload-time = "2024-10-25T17:25:40.039Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c1/8b/5fe2cc11fee489817272089c4203e679c63b570a5aaeb18d852ae3cbba6a/et_xmlfile-2.0.0-py3-none-any.whl", hash = "sha256:7a91720bc756843502c3b7504c77b8fe44217c85c537d85037f0f536151b2caa", size = 18059, upload-time = "2024-10-25T17:25:39.051Z" }, +] + +[[package]] +name = "filelock" +version = "3.16.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/9d/db/3ef5bb276dae18d6ec2124224403d1d67bccdbefc17af4cc8f553e341ab1/filelock-3.16.1.tar.gz", hash = "sha256:c249fbfcd5db47e5e2d6d62198e565475ee65e4831e2561c8e313fa7eb961435", size = 18037, upload-time = "2024-09-17T19:02:01.779Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b9/f8/feced7779d755758a52d1f6635d990b8d98dc0a29fa568bbe0625f18fdf3/filelock-3.16.1-py3-none-any.whl", hash = "sha256:2082e5703d51fbf98ea75855d9d5527e33d8ff23099bec374a134febee6946b0", size = 16163, upload-time = "2024-09-17T19:02:00.268Z" }, +] + +[[package]] +name = "filelock" +version = "3.19.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.13.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'darwin')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'darwin')", + "(python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform == 'darwin')", + "(python_full_version == '3.10.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'darwin')", + "(python_full_version == '3.9.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform == 'darwin')", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/40/bb/0ab3e58d22305b6f5440629d20683af28959bf793d98d11950e305c1c326/filelock-3.19.1.tar.gz", hash = "sha256:66eda1888b0171c998b35be2bcc0f6d75c388a7ce20c3f3f37aa8e96c2dddf58", size = 17687, upload-time = "2025-08-14T16:56:03.016Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/42/14/42b2651a2f46b022ccd948bca9f2d5af0fd8929c4eec235b8d6d844fbe67/filelock-3.19.1-py3-none-any.whl", hash = "sha256:d38e30481def20772f5baf097c122c3babc4fcdb7e14e57049eb9d88c6dc017d", size = 15988, upload-time = "2025-08-14T16:56:01.633Z" }, +] + +[[package]] +name = "fsspec" +version = "2025.3.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/34/f4/5721faf47b8c499e776bc34c6a8fc17efdf7fdef0b00f398128bc5dcb4ac/fsspec-2025.3.0.tar.gz", hash = "sha256:a935fd1ea872591f2b5148907d103488fc523295e6c64b835cfad8c3eca44972", size = 298491, upload-time = "2025-03-07T21:47:56.461Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/56/53/eb690efa8513166adef3e0669afd31e95ffde69fb3c52ec2ac7223ed6018/fsspec-2025.3.0-py3-none-any.whl", hash = "sha256:efb87af3efa9103f94ca91a7f8cb7a4df91af9f74fc106c9c7ea0efd7277c1b3", size = 193615, upload-time = "2025-03-07T21:47:54.809Z" }, +] + +[[package]] +name = "fsspec" +version = "2025.7.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.13.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'darwin')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'darwin')", + "(python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform == 'darwin')", + "(python_full_version == '3.10.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'darwin')", + "(python_full_version == '3.9.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform == 'darwin')", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8b/02/0835e6ab9cfc03916fe3f78c0956cfcdb6ff2669ffa6651065d5ebf7fc98/fsspec-2025.7.0.tar.gz", hash = "sha256:786120687ffa54b8283d942929540d8bc5ccfa820deb555a2b5d0ed2b737bf58", size = 304432, upload-time = "2025-07-15T16:05:21.19Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2f/e0/014d5d9d7a4564cf1c40b5039bc882db69fd881111e03ab3657ac0b218e2/fsspec-2025.7.0-py3-none-any.whl", hash = "sha256:8b012e39f63c7d5f10474de957f3ab793b47b45ae7d39f2fb735f8bbe25c0e21", size = 199597, upload-time = "2025-07-15T16:05:19.529Z" }, +] + +[[package]] +name = "funasr-bmodel" +version = "0.1.0" +source = { virtual = "." } +dependencies = [ + { name = "hydra-core" }, + { name = "kaldiio" }, + { name = "librosa" }, + { name = "omegaconf" }, + { name = "rotary-embedding-torch" }, + { name = "six" }, + { name = "torch-complex" }, + { name = "torchaudio", version = "2.5.1+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "python_full_version < '3.9'" }, + { name = "torchaudio", version = "2.8.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(python_full_version >= '3.9' and python_full_version < '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.9' and sys_platform == 'darwin')" }, + { name = "torchaudio", version = "2.8.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.9' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.9' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.9' and sys_platform != 'darwin' and sys_platform != 'linux')" }, + { name = "tpu-perf" }, + { name = "tqdm" }, + { name = "umap-learn", version = "0.5.7", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "umap-learn", version = "0.5.9.post2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9'" }, +] + +[package.metadata] +requires-dist = [ + { name = "hydra-core", specifier = ">=1.3.2" }, + { name = "kaldiio", specifier = ">=2.18.1" }, + { name = "librosa", specifier = ">=0.11.0" }, + { name = "omegaconf", specifier = ">=2.3.0" }, + { name = "rotary-embedding-torch", specifier = ">=0.8.9" }, + { name = "six", specifier = ">=1.17.0" }, + { name = "torch-complex", specifier = ">=0.4.4" }, + { name = "torchaudio", specifier = ">=2.5.1", index = "https://download.pytorch.org/whl/cpu" }, + { name = "tpu-perf", path = "bmodel/tpu_perf-1.2.35-py3-none-manylinux2014_aarch64.whl" }, + { name = "tqdm", specifier = ">=4.67.1" }, + { name = "umap-learn", specifier = ">=0.5.7" }, +] + +[[package]] +name = "hydra-core" +version = "1.3.2" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +dependencies = [ + { name = "antlr4-python3-runtime" }, + { name = "importlib-resources", marker = "python_full_version < '3.9'" }, + { name = "omegaconf" }, + { name = "packaging" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6d/8e/07e42bc434a847154083b315779b0a81d567154504624e181caf2c71cd98/hydra-core-1.3.2.tar.gz", hash = "sha256:8a878ed67216997c3e9d88a8e72e7b4767e81af37afb4ea3334b269a4390a824", size = 3263494, upload-time = "2023-02-23T18:33:43.03Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c6/50/e0edd38dcd63fb26a8547f13d28f7a008bc4a3fd4eb4ff030673f22ad41a/hydra_core-1.3.2-py3-none-any.whl", hash = "sha256:fa0238a9e31df3373b35b0bfb672c34cc92718d21f81311d8996a16de1141d8b", size = 154547, upload-time = "2023-02-23T18:33:40.801Z" }, +] + +[[package]] +name = "idna" +version = "3.10" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f1/70/7703c29685631f5a7590aa73f1f1d3fa9a380e654b86af429e0934a32f7d/idna-3.10.tar.gz", hash = "sha256:12f65c9b470abda6dc35cf8e63cc574b1c52b11df2c86030af0ac09b01b13ea9", size = 190490, upload-time = "2024-09-15T18:07:39.745Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/76/c6/c88e154df9c4e1a2a66ccf0005a88dfb2650c1dffb6f5ce603dfbd452ce3/idna-3.10-py3-none-any.whl", hash = "sha256:946d195a0d259cbba61165e88e65941f16e9b36ea6ddb97f00452bae8b1287d3", size = 70442, upload-time = "2024-09-15T18:07:37.964Z" }, +] + +[[package]] +name = "importlib-metadata" +version = "8.5.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] +dependencies = [ + { name = "zipp", version = "3.20.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/cd/12/33e59336dca5be0c398a7482335911a33aa0e20776128f038019f1a95f1b/importlib_metadata-8.5.0.tar.gz", hash = "sha256:71522656f0abace1d072b9e5481a48f07c138e00f079c38c8f883823f9c26bd7", size = 55304, upload-time = "2024-09-11T14:56:08.937Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a0/d9/a1e041c5e7caa9a05c925f4bdbdfb7f006d1f74996af53467bc394c97be7/importlib_metadata-8.5.0-py3-none-any.whl", hash = "sha256:45e54197d28b7a7f1559e60b95e7c567032b602131fbd588f1497f47880aa68b", size = 26514, upload-time = "2024-09-11T14:56:07.019Z" }, +] + +[[package]] +name = "importlib-metadata" +version = "8.7.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", +] +dependencies = [ + { name = "zipp", version = "3.23.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/76/66/650a33bd90f786193e4de4b3ad86ea60b53c89b669a5c7be931fac31cdb0/importlib_metadata-8.7.0.tar.gz", hash = "sha256:d13b81ad223b890aa16c5471f2ac3056cf76c5f10f82d6f9292f0b415f389000", size = 56641, upload-time = "2025-04-27T15:29:01.736Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/20/b0/36bd937216ec521246249be3bf9855081de4c5e06a0c9b4219dbeda50373/importlib_metadata-8.7.0-py3-none-any.whl", hash = "sha256:e5dd1551894c77868a30651cef00984d50e1002d06942a7101d34870c5f02afd", size = 27656, upload-time = "2025-04-27T15:29:00.214Z" }, +] + +[[package]] +name = "importlib-resources" +version = "6.4.5" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +dependencies = [ + { name = "zipp", version = "3.20.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/98/be/f3e8c6081b684f176b761e6a2fef02a0be939740ed6f54109a2951d806f3/importlib_resources-6.4.5.tar.gz", hash = "sha256:980862a1d16c9e147a59603677fa2aa5fd82b87f223b6cb870695bcfce830065", size = 43372, upload-time = "2024-09-09T17:03:14.677Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e1/6a/4604f9ae2fa62ef47b9de2fa5ad599589d28c9fd1d335f32759813dfa91e/importlib_resources-6.4.5-py3-none-any.whl", hash = "sha256:ac29d5f956f01d5e4bb63102a5a19957f1b9175e45649977264a1416783bb717", size = 36115, upload-time = "2024-09-09T17:03:13.39Z" }, +] + +[[package]] +name = "jinja2" +version = "3.1.6" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +dependencies = [ + { name = "markupsafe", version = "2.1.5", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "markupsafe", version = "3.0.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115, upload-time = "2025-03-05T20:05:02.478Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, +] + +[[package]] +name = "joblib" +version = "1.4.2" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/64/33/60135848598c076ce4b231e1b1895170f45fbcaeaa2c9d5e38b04db70c35/joblib-1.4.2.tar.gz", hash = "sha256:2382c5816b2636fbd20a09e0f4e9dad4736765fdfb7dca582943b9c1366b3f0e", size = 2116621, upload-time = "2024-05-02T12:15:05.765Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/91/29/df4b9b42f2be0b623cbd5e2140cafcaa2bef0759a00b7b70104dcfe2fb51/joblib-1.4.2-py3-none-any.whl", hash = "sha256:06d478d5674cbc267e7496a410ee875abd68e4340feff4490bcb7afb88060ae6", size = 301817, upload-time = "2024-05-02T12:15:00.765Z" }, +] + +[[package]] +name = "joblib" +version = "1.5.2" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.13.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'darwin')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'darwin')", + "(python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform == 'darwin')", + "(python_full_version == '3.10.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'darwin')", + "(python_full_version == '3.9.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform == 'darwin')", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e8/5d/447af5ea094b9e4c4054f82e223ada074c552335b9b4b2d14bd9b35a67c4/joblib-1.5.2.tar.gz", hash = "sha256:3faa5c39054b2f03ca547da9b2f52fde67c06240c31853f306aea97f13647b55", size = 331077, upload-time = "2025-08-27T12:15:46.575Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1e/e8/685f47e0d754320684db4425a0967f7d3fa70126bffd76110b7009a0090f/joblib-1.5.2-py3-none-any.whl", hash = "sha256:4e1f0bdbb987e6d843c70cf43714cb276623def372df3c22fe5266b2670bc241", size = 308396, upload-time = "2025-08-27T12:15:45.188Z" }, +] + +[[package]] +name = "kaldiio" +version = "2.18.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +dependencies = [ + { name = "numpy", version = "1.24.4", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "numpy", version = "2.0.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.10'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8d/85/92435e8e62eb3d43eded9f24643fc2a6dbce031cebceed11528147c7873f/kaldiio-2.18.1.tar.gz", hash = "sha256:0283d197fac6ac683f7a9e6af8d18aad9dbd2c4a997f22e45294f2ac1ee3c432", size = 35570, upload-time = "2025-03-06T15:57:52.375Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ae/e3/6c3b42233225f398f7a72988b524f654ae818cca0d441db847a2761203e9/kaldiio-2.18.1-py3-none-any.whl", hash = "sha256:397a4cd18977acaae7acabfba6807ee0a6978c620064381a266eac15b3c1a0a0", size = 29330, upload-time = "2025-03-06T15:57:50.82Z" }, +] + +[[package]] +name = "lazy-loader" +version = "0.4" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +dependencies = [ + { name = "packaging" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6f/6b/c875b30a1ba490860c93da4cabf479e03f584eba06fe5963f6f6644653d8/lazy_loader-0.4.tar.gz", hash = "sha256:47c75182589b91a4e1a85a136c074285a5ad4d9f39c63e0d7fb76391c4574cd1", size = 15431, upload-time = "2024-04-05T13:03:12.261Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/83/60/d497a310bde3f01cb805196ac61b7ad6dc5dcf8dce66634dc34364b20b4f/lazy_loader-0.4-py3-none-any.whl", hash = "sha256:342aa8e14d543a154047afb4ba8ef17f5563baad3fc610d7b15b213b0f119efc", size = 12097, upload-time = "2024-04-05T13:03:10.514Z" }, +] + +[[package]] +name = "librosa" +version = "0.11.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +dependencies = [ + { name = "audioread" }, + { name = "decorator" }, + { name = "joblib", version = "1.4.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "joblib", version = "1.5.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9'" }, + { name = "lazy-loader" }, + { name = "msgpack" }, + { name = "numba", version = "0.58.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "numba", version = "0.60.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, + { name = "numba", version = "0.61.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.10'" }, + { name = "numpy", version = "1.24.4", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "numpy", version = "2.0.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.10'" }, + { name = "pooch" }, + { name = "scikit-learn", version = "1.3.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "scikit-learn", version = "1.6.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, + { name = "scikit-learn", version = "1.7.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.10'" }, + { name = "scipy", version = "1.10.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "scipy", version = "1.13.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, + { name = "scipy", version = "1.15.3", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.10.*'" }, + { name = "scipy", version = "1.16.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "soundfile" }, + { name = "soxr", version = "0.3.7", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "soxr", version = "0.5.0.post1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9'" }, + { name = "standard-aifc", marker = "python_full_version >= '3.13'" }, + { name = "standard-sunau", marker = "python_full_version >= '3.13'" }, + { name = "typing-extensions", version = "4.13.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "typing-extensions", version = "4.15.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/64/36/360b5aafa0238e29758729e9486c6ed92a6f37fa403b7875e06c115cdf4a/librosa-0.11.0.tar.gz", hash = "sha256:f5ed951ca189b375bbe2e33b2abd7e040ceeee302b9bbaeeffdfddb8d0ace908", size = 327001, upload-time = "2025-03-11T15:09:54.884Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b5/ba/c63c5786dfee4c3417094c4b00966e61e4a63efecee22cb7b4c0387dda83/librosa-0.11.0-py3-none-any.whl", hash = "sha256:0b6415c4fd68bff4c29288abe67c6d80b587e0e1e2cfb0aad23e4559504a7fa1", size = 260749, upload-time = "2025-03-11T15:09:52.982Z" }, +] + +[[package]] +name = "llvmlite" +version = "0.41.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/01/c6/bc6634da9f58edf91a1a002280c6380f404715245a49a46234b1d9d9585a/llvmlite-0.41.1.tar.gz", hash = "sha256:f19f767a018e6ec89608e1f6b13348fa2fcde657151137cb64e56d48598a92db", size = 146564, upload-time = "2023-10-18T13:36:19.997Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/10/e2/931d2d35350d3c375d6c43a04982fb5d8f90c15002de024a814d919a1e67/llvmlite-0.41.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:c1e1029d47ee66d3a0c4d6088641882f75b93db82bd0e6178f7bd744ebce42b9", size = 31032708, upload-time = "2023-10-18T13:34:42.554Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/47/d5/2d6261a47dd6b1005d60a6e37a3de007ddc15a297f146ee905605e3cbc3e/llvmlite-0.41.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:150d0bc275a8ac664a705135e639178883293cf08c1a38de3bbaa2f693a0a867", size = 28753858, upload-time = "2023-10-18T13:34:47.195Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/da/5e/afce5c757df730be7706ae671f5fce2cacdf4e46eed2b1a916e595453388/llvmlite-0.41.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1eee5cf17ec2b4198b509272cf300ee6577229d237c98cc6e63861b08463ddc6", size = 42623412, upload-time = "2023-10-18T13:34:54.423Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/57/7d/ef28d5812f852b93bd2a583d00cdcde56833d31b645ae0eaa7e71eecfb4e/llvmlite-0.41.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0dd0338da625346538f1173a17cabf21d1e315cf387ca21b294ff209d176e244", size = 43623334, upload-time = "2023-10-18T13:34:58.57Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/48/39/7db84e829b874f90316e9f6d7497f761b9a645abe61f8c9588b292876314/llvmlite-0.41.1-cp310-cp310-win32.whl", hash = "sha256:fa1469901a2e100c17eb8fe2678e34bd4255a3576d1a543421356e9c14d6e2ae", size = 23414330, upload-time = "2023-10-18T13:35:03.791Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1f/05/cfc9c51aff9adbb6f7df93df3fb20100ee5d9c3d0445ade3df8057e50b0a/llvmlite-0.41.1-cp310-cp310-win_amd64.whl", hash = "sha256:2b76acee82ea0e9304be6be9d4b3840208d050ea0dcad75b1635fa06e949a0ae", size = 28108282, upload-time = "2023-10-18T13:35:07.674Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/96/e0/54b0a58969b208b3a66ed30f560e127f52c3add6027b8e67d0a00ea0b178/llvmlite-0.41.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:210e458723436b2469d61b54b453474e09e12a94453c97ea3fbb0742ba5a83d8", size = 31032706, upload-time = "2023-10-18T13:35:11.181Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/12/f4/9bc78ef5d52cbe595a849d332c872583a6294afff6f5317de6dce4abbe55/llvmlite-0.41.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:855f280e781d49e0640aef4c4af586831ade8f1a6c4df483fb901cbe1a48d127", size = 28753841, upload-time = "2023-10-18T13:35:14.601Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b5/b0/48caac8ea81d6a11a1917dd98258ea67c9b8a78a31fa45b89ae6aa94e678/llvmlite-0.41.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b67340c62c93a11fae482910dc29163a50dff3dfa88bc874872d28ee604a83be", size = 42623411, upload-time = "2023-10-18T13:35:19.577Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0a/dd/e10658cb0e4ad8364b37ae2a6187ebb6b3b35e08b36ef9e98790b9f2f437/llvmlite-0.41.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2181bb63ef3c607e6403813421b46982c3ac6bfc1f11fa16a13eaafb46f578e6", size = 43623334, upload-time = "2023-10-18T13:35:24.741Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/14/fe/d3a9c921a2adad2e9f24693754983f290e0dac9410666e802b9dba4d0218/llvmlite-0.41.1-cp311-cp311-win_amd64.whl", hash = "sha256:9564c19b31a0434f01d2025b06b44c7ed422f51e719ab5d24ff03b7560066c9a", size = 28092676, upload-time = "2023-10-18T13:35:28.992Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/04/ad/f5d3f8accf618e136a205de92db49516d1e9a868415164dab407107cece5/llvmlite-0.41.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:5940bc901fb0325970415dbede82c0b7f3e35c2d5fd1d5e0047134c2c46b3281", size = 31032702, upload-time = "2023-10-18T13:35:32.88Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a7/07/c5c053cebc0d16fbbe5329fd10b9d141681fba7a805796e8517a1217d4e8/llvmlite-0.41.1-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:8b0a9a47c28f67a269bb62f6256e63cef28d3c5f13cbae4fab587c3ad506778b", size = 28753852, upload-time = "2023-10-18T13:35:36.198Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/37/c0/2fcbb511e4af6f4bbd053b1ce16ae216ff088f3255add6b516c34175bfbc/llvmlite-0.41.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f8afdfa6da33f0b4226af8e64cfc2b28986e005528fbf944d0a24a72acfc9432", size = 42623411, upload-time = "2023-10-18T13:35:39.862Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/46/a5/de824aba87a62e0da43f9b590ac91709d460a61e03f1b430b83ca64b76dc/llvmlite-0.41.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8454c1133ef701e8c050a59edd85d238ee18bb9a0eb95faf2fca8b909ee3c89a", size = 43623331, upload-time = "2023-10-18T13:35:44.69Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7e/f0/066eb0d43d276da3fc39cf61ad956293a2d3c97f4e976fcc58d3519862e8/llvmlite-0.41.1-cp38-cp38-win32.whl", hash = "sha256:2d92c51e6e9394d503033ffe3292f5bef1566ab73029ec853861f60ad5c925d0", size = 23414330, upload-time = "2023-10-18T13:35:48.729Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1b/57/6da776f73951956ece7d8fda605c69d512a6724bb7b255adf9d48ef7941d/llvmlite-0.41.1-cp38-cp38-win_amd64.whl", hash = "sha256:df75594e5a4702b032684d5481db3af990b69c249ccb1d32687b8501f0689432", size = 28108279, upload-time = "2023-10-18T13:35:51.863Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fc/7a/c6741000d767fed4b339fcd4fa65afbc5fe776473d5f9e9c41eceab0a7c6/llvmlite-0.41.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:04725975e5b2af416d685ea0769f4ecc33f97be541e301054c9f741003085802", size = 31032700, upload-time = "2023-10-18T13:35:55.347Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8b/66/b3edcd1b2631af2daaf7a09520c5ff863e459eb17dd28f7d4cac63f3532c/llvmlite-0.41.1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:bf14aa0eb22b58c231243dccf7e7f42f7beec48970f2549b3a6acc737d1a4ba4", size = 28753838, upload-time = "2023-10-18T13:35:58.786Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d3/5b/5af7cc0d2097ada6372b8b90986b45267b7d137a4fda58d7534d867572b2/llvmlite-0.41.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:92c32356f669e036eb01016e883b22add883c60739bc1ebee3a1cc0249a50828", size = 42623409, upload-time = "2023-10-18T13:36:02.781Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/be/a6/55020df7f04666bceb3ccc5dc0cb0812d68d5fdec883ff015a9110928829/llvmlite-0.41.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:24091a6b31242bcdd56ae2dbea40007f462260bc9bdf947953acc39dffd54f8f", size = 43623330, upload-time = "2023-10-18T13:36:09.753Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/07/73/0aef09cf017a4f8b9c0eb371d273e7bb3e514992c4f1d8ed09e8f81f4602/llvmlite-0.41.1-cp39-cp39-win32.whl", hash = "sha256:880cb57ca49e862e1cd077104375b9d1dfdc0622596dfa22105f470d7bacb309", size = 23414330, upload-time = "2023-10-18T13:36:13.771Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/88/99/aa829162a871f0f91dba9a0344726449cafbff4e9e55d7a714051cbb6215/llvmlite-0.41.1-cp39-cp39-win_amd64.whl", hash = "sha256:92f093986ab92e71c9ffe334c002f96defc7986efda18397d0f08534f3ebdc4d", size = 28092675, upload-time = "2023-10-18T13:36:17.19Z" }, +] + +[[package]] +name = "llvmlite" +version = "0.43.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform == 'darwin')", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/9f/3d/f513755f285db51ab363a53e898b85562e950f79a2e6767a364530c2f645/llvmlite-0.43.0.tar.gz", hash = "sha256:ae2b5b5c3ef67354824fb75517c8db5fbe93bc02cd9671f3c62271626bc041d5", size = 157069, upload-time = "2024-06-13T18:09:32.641Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/23/ff/6ca7e98998b573b4bd6566f15c35e5c8bea829663a6df0c7aa55ab559da9/llvmlite-0.43.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:a289af9a1687c6cf463478f0fa8e8aa3b6fb813317b0d70bf1ed0759eab6f761", size = 31064408, upload-time = "2024-06-13T18:08:13.462Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ca/5c/a27f9257f86f0cda3f764ff21d9f4217b9f6a0d45e7a39ecfa7905f524ce/llvmlite-0.43.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:6d4fd101f571a31acb1559ae1af30f30b1dc4b3186669f92ad780e17c81e91bc", size = 28793153, upload-time = "2024-06-13T18:08:17.336Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7e/3c/4410f670ad0a911227ea2ecfcba9f672a77cf1924df5280c4562032ec32d/llvmlite-0.43.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7d434ec7e2ce3cc8f452d1cd9a28591745de022f931d67be688a737320dfcead", size = 42857276, upload-time = "2024-06-13T18:08:21.071Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c6/21/2ffbab5714e72f2483207b4a1de79b2eecd9debbf666ff4e7067bcc5c134/llvmlite-0.43.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6912a87782acdff6eb8bf01675ed01d60ca1f2551f8176a300a886f09e836a6a", size = 43871781, upload-time = "2024-06-13T18:08:26.32Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f2/26/b5478037c453554a61625ef1125f7e12bb1429ae11c6376f47beba9b0179/llvmlite-0.43.0-cp310-cp310-win_amd64.whl", hash = "sha256:14f0e4bf2fd2d9a75a3534111e8ebeb08eda2f33e9bdd6dfa13282afacdde0ed", size = 28123487, upload-time = "2024-06-13T18:08:30.348Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/95/8c/de3276d773ab6ce3ad676df5fab5aac19696b2956319d65d7dd88fb10f19/llvmlite-0.43.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:3e8d0618cb9bfe40ac38a9633f2493d4d4e9fcc2f438d39a4e854f39cc0f5f98", size = 31064409, upload-time = "2024-06-13T18:08:34.006Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ee/e1/38deed89ced4cf378c61e232265cfe933ccde56ae83c901aa68b477d14b1/llvmlite-0.43.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:e0a9a1a39d4bf3517f2af9d23d479b4175ead205c592ceeb8b89af48a327ea57", size = 28793149, upload-time = "2024-06-13T18:08:37.42Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2f/b2/4429433eb2dc8379e2cb582502dca074c23837f8fd009907f78a24de4c25/llvmlite-0.43.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c1da416ab53e4f7f3bc8d4eeba36d801cc1894b9fbfbf2022b29b6bad34a7df2", size = 42857277, upload-time = "2024-06-13T18:08:40.822Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6b/99/5d00a7d671b1ba1751fc9f19d3b36f3300774c6eebe2bcdb5f6191763eb4/llvmlite-0.43.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:977525a1e5f4059316b183fb4fd34fa858c9eade31f165427a3977c95e3ee749", size = 43871781, upload-time = "2024-06-13T18:08:46.41Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/20/ab/ed5ed3688c6ba4f0b8d789da19fd8e30a9cf7fc5852effe311bc5aefe73e/llvmlite-0.43.0-cp311-cp311-win_amd64.whl", hash = "sha256:d5bd550001d26450bd90777736c69d68c487d17bf371438f975229b2b8241a91", size = 28107433, upload-time = "2024-06-13T18:08:50.834Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0b/67/9443509e5d2b6d8587bae3ede5598fa8bd586b1c7701696663ea8af15b5b/llvmlite-0.43.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:f99b600aa7f65235a5a05d0b9a9f31150c390f31261f2a0ba678e26823ec38f7", size = 31064409, upload-time = "2024-06-13T18:08:54.375Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a2/9c/24139d3712d2d352e300c39c0e00d167472c08b3bd350c3c33d72c88ff8d/llvmlite-0.43.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:35d80d61d0cda2d767f72de99450766250560399edc309da16937b93d3b676e7", size = 28793145, upload-time = "2024-06-13T18:08:57.953Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/bf/f1/4c205a48488e574ee9f6505d50e84370a978c90f08dab41a42d8f2c576b6/llvmlite-0.43.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:eccce86bba940bae0d8d48ed925f21dbb813519169246e2ab292b5092aba121f", size = 42857276, upload-time = "2024-06-13T18:09:02.067Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/00/5f/323c4d56e8401c50185fd0e875fcf06b71bf825a863699be1eb10aa2a9cb/llvmlite-0.43.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:df6509e1507ca0760787a199d19439cc887bfd82226f5af746d6977bd9f66844", size = 43871781, upload-time = "2024-06-13T18:09:06.667Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c6/94/dea10e263655ce78d777e78d904903faae39d1fc440762be4a9dc46bed49/llvmlite-0.43.0-cp312-cp312-win_amd64.whl", hash = "sha256:7a2872ee80dcf6b5dbdc838763d26554c2a18aa833d31a2635bff16aafefb9c9", size = 28107442, upload-time = "2024-06-13T18:09:10.709Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2a/73/12925b1bbb3c2beb6d96f892ef5b4d742c34f00ddb9f4a125e9e87b22f52/llvmlite-0.43.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:9cd2a7376f7b3367019b664c21f0c61766219faa3b03731113ead75107f3b66c", size = 31064410, upload-time = "2024-06-13T18:09:14.091Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/cc/61/58c70aa0808a8cba825a7d98cc65bef4801b99328fba80837bfcb5fc767f/llvmlite-0.43.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:18e9953c748b105668487b7c81a3e97b046d8abf95c4ddc0cd3c94f4e4651ae8", size = 28793145, upload-time = "2024-06-13T18:09:17.531Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c8/c6/9324eb5de2ba9d99cbed853d85ba7a318652a48e077797bec27cf40f911d/llvmlite-0.43.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:74937acd22dc11b33946b67dca7680e6d103d6e90eeaaaf932603bec6fe7b03a", size = 42857276, upload-time = "2024-06-13T18:09:21.377Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e0/d0/889e9705107db7b1ec0767b03f15d7b95b4c4f9fdf91928ab1c7e9ffacf6/llvmlite-0.43.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bc9efc739cc6ed760f795806f67889923f7274276f0eb45092a1473e40d9b867", size = 43871777, upload-time = "2024-06-13T18:09:25.76Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/df/41/73cc26a2634b538cfe813f618c91e7e9960b8c163f8f0c94a2b0f008b9da/llvmlite-0.43.0-cp39-cp39-win_amd64.whl", hash = "sha256:47e147cdda9037f94b399bf03bfd8a6b6b1f2f90be94a454e3386f006455a9b4", size = 28123489, upload-time = "2024-06-13T18:09:29.78Z" }, +] + +[[package]] +name = "llvmlite" +version = "0.44.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.13.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'darwin')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'darwin')", + "(python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform == 'darwin')", + "(python_full_version == '3.10.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'darwin')", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/89/6a/95a3d3610d5c75293d5dbbb2a76480d5d4eeba641557b69fe90af6c5b84e/llvmlite-0.44.0.tar.gz", hash = "sha256:07667d66a5d150abed9157ab6c0b9393c9356f229784a4385c02f99e94fc94d4", size = 171880, upload-time = "2025-01-20T11:14:41.342Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/41/75/d4863ddfd8ab5f6e70f4504cf8cc37f4e986ec6910f4ef8502bb7d3c1c71/llvmlite-0.44.0-cp310-cp310-macosx_10_14_x86_64.whl", hash = "sha256:9fbadbfba8422123bab5535b293da1cf72f9f478a65645ecd73e781f962ca614", size = 28132306, upload-time = "2025-01-20T11:12:18.634Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/37/d9/6e8943e1515d2f1003e8278819ec03e4e653e2eeb71e4d00de6cfe59424e/llvmlite-0.44.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:cccf8eb28f24840f2689fb1a45f9c0f7e582dd24e088dcf96e424834af11f791", size = 26201096, upload-time = "2025-01-20T11:12:24.544Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/aa/46/8ffbc114def88cc698906bf5acab54ca9fdf9214fe04aed0e71731fb3688/llvmlite-0.44.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7202b678cdf904823c764ee0fe2dfe38a76981f4c1e51715b4cb5abb6cf1d9e8", size = 42361859, upload-time = "2025-01-20T11:12:31.839Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/30/1c/9366b29ab050a726af13ebaae8d0dff00c3c58562261c79c635ad4f5eb71/llvmlite-0.44.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:40526fb5e313d7b96bda4cbb2c85cd5374e04d80732dd36a282d72a560bb6408", size = 41184199, upload-time = "2025-01-20T11:12:40.049Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/69/07/35e7c594b021ecb1938540f5bce543ddd8713cff97f71d81f021221edc1b/llvmlite-0.44.0-cp310-cp310-win_amd64.whl", hash = "sha256:41e3839150db4330e1b2716c0be3b5c4672525b4c9005e17c7597f835f351ce2", size = 30332381, upload-time = "2025-01-20T11:12:47.054Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b5/e2/86b245397052386595ad726f9742e5223d7aea999b18c518a50e96c3aca4/llvmlite-0.44.0-cp311-cp311-macosx_10_14_x86_64.whl", hash = "sha256:eed7d5f29136bda63b6d7804c279e2b72e08c952b7c5df61f45db408e0ee52f3", size = 28132305, upload-time = "2025-01-20T11:12:53.936Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ff/ec/506902dc6870249fbe2466d9cf66d531265d0f3a1157213c8f986250c033/llvmlite-0.44.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:ace564d9fa44bb91eb6e6d8e7754977783c68e90a471ea7ce913bff30bd62427", size = 26201090, upload-time = "2025-01-20T11:12:59.847Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/99/fe/d030f1849ebb1f394bb3f7adad5e729b634fb100515594aca25c354ffc62/llvmlite-0.44.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c5d22c3bfc842668168a786af4205ec8e3ad29fb1bc03fd11fd48460d0df64c1", size = 42361858, upload-time = "2025-01-20T11:13:07.623Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d7/7a/ce6174664b9077fc673d172e4c888cb0b128e707e306bc33fff8c2035f0d/llvmlite-0.44.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f01a394e9c9b7b1d4e63c327b096d10f6f0ed149ef53d38a09b3749dcf8c9610", size = 41184200, upload-time = "2025-01-20T11:13:20.058Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5f/c6/258801143975a6d09a373f2641237992496e15567b907a4d401839d671b8/llvmlite-0.44.0-cp311-cp311-win_amd64.whl", hash = "sha256:d8489634d43c20cd0ad71330dde1d5bc7b9966937a263ff1ec1cebb90dc50955", size = 30331193, upload-time = "2025-01-20T11:13:26.976Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/15/86/e3c3195b92e6e492458f16d233e58a1a812aa2bfbef9bdd0fbafcec85c60/llvmlite-0.44.0-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:1d671a56acf725bf1b531d5ef76b86660a5ab8ef19bb6a46064a705c6ca80aad", size = 28132297, upload-time = "2025-01-20T11:13:32.57Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d6/53/373b6b8be67b9221d12b24125fd0ec56b1078b660eeae266ec388a6ac9a0/llvmlite-0.44.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:5f79a728e0435493611c9f405168682bb75ffd1fbe6fc360733b850c80a026db", size = 26201105, upload-time = "2025-01-20T11:13:38.744Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/cb/da/8341fd3056419441286c8e26bf436923021005ece0bff5f41906476ae514/llvmlite-0.44.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c0143a5ef336da14deaa8ec26c5449ad5b6a2b564df82fcef4be040b9cacfea9", size = 42361901, upload-time = "2025-01-20T11:13:46.711Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/53/ad/d79349dc07b8a395a99153d7ce8b01d6fcdc9f8231355a5df55ded649b61/llvmlite-0.44.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d752f89e31b66db6f8da06df8b39f9b91e78c5feea1bf9e8c1fba1d1c24c065d", size = 41184247, upload-time = "2025-01-20T11:13:56.159Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e2/3b/a9a17366af80127bd09decbe2a54d8974b6d8b274b39bf47fbaedeec6307/llvmlite-0.44.0-cp312-cp312-win_amd64.whl", hash = "sha256:eae7e2d4ca8f88f89d315b48c6b741dcb925d6a1042da694aa16ab3dd4cbd3a1", size = 30332380, upload-time = "2025-01-20T11:14:02.442Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/89/24/4c0ca705a717514c2092b18476e7a12c74d34d875e05e4d742618ebbf449/llvmlite-0.44.0-cp313-cp313-macosx_10_14_x86_64.whl", hash = "sha256:319bddd44e5f71ae2689859b7203080716448a3cd1128fb144fe5c055219d516", size = 28132306, upload-time = "2025-01-20T11:14:09.035Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/01/cf/1dd5a60ba6aee7122ab9243fd614abcf22f36b0437cbbe1ccf1e3391461c/llvmlite-0.44.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:9c58867118bad04a0bb22a2e0068c693719658105e40009ffe95c7000fcde88e", size = 26201090, upload-time = "2025-01-20T11:14:15.401Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d2/1b/656f5a357de7135a3777bd735cc7c9b8f23b4d37465505bd0eaf4be9befe/llvmlite-0.44.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:46224058b13c96af1365290bdfebe9a6264ae62fb79b2b55693deed11657a8bf", size = 42361904, upload-time = "2025-01-20T11:14:22.949Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d8/e1/12c5f20cb9168fb3464a34310411d5ad86e4163c8ff2d14a2b57e5cc6bac/llvmlite-0.44.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:aa0097052c32bf721a4efc03bd109d335dfa57d9bffb3d4c24cc680711b8b4fc", size = 41184245, upload-time = "2025-01-20T11:14:31.731Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d0/81/e66fc86539293282fd9cb7c9417438e897f369e79ffb62e1ae5e5154d4dd/llvmlite-0.44.0-cp313-cp313-win_amd64.whl", hash = "sha256:2fb7c4f2fb86cbae6dca3db9ab203eeea0e22d73b99bc2341cdf9de93612e930", size = 30331193, upload-time = "2025-01-20T11:14:38.578Z" }, +] + +[[package]] +name = "lmdb" +version = "1.7.3" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5e/87/b6f8bccab1ec1fe58866108c0da648715bc1488627edf9702dae6380613a/lmdb-1.7.3.tar.gz", hash = "sha256:d4a27b7af4fe38f3409d9fbfbf47b617b3d94dee98a00bc8c2a8336ccd3f9b15", size = 883449, upload-time = "2025-07-26T01:11:43.197Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/39/82/fd14062511f697f469d34ea8b87575176a96d2798606fe9efc5ca3a40f4b/lmdb-1.7.3-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:d604665b79eb55b41fb8eb11da12438e59efcf1f3e27aeabd7a112c10932f17a", size = 101028, upload-time = "2025-07-26T01:10:47.853Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/84/56/cb61d542d826beccbce93d400ebbd16312ddffe1195531dc7c05e288bc67/lmdb-1.7.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:19b6c726d69f163fe413befa2daa5b3cfb3b9f18b2aac99b3af40b70b7b64dd7", size = 98456, upload-time = "2025-07-26T01:10:49.455Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b7/b7/7e85c30bb50ba93cf77a9c137f25dbb2857a22d8004dc22607d3add728ee/lmdb-1.7.3-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:94ba8c056f001ea138e478a9faef0c3a54e84d6f2525b1e6ef771563e143e682", size = 295196, upload-time = "2025-07-26T01:10:50.79Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/83/76/1800f3605a8b94462be1e4c62c046fb8a93a6b1ef7eeb4e8ef4f5a0b36e0/lmdb-1.7.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1822c07a2a090b283e8ff68e3c83e361a99872fb6f2da1b96f2cbee33cc59763", size = 296560, upload-time = "2025-07-26T01:10:52.032Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e3/e6/cc155fc4a81ee5f1860df425ea2404f5a092eaa7a37d1298c463cada5d6a/lmdb-1.7.3-cp310-cp310-win_amd64.whl", hash = "sha256:11b62fd18d9bc41c5a8270453d76e317dec4cd744ba0a6ce65807a13bdd02d6d", size = 99267, upload-time = "2025-07-26T01:10:53.351Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/62/b0/9910c56d64e3b1a0331235762ff192446b63aa207c5e2bf6c7a9a9d08ddc/lmdb-1.7.3-cp310-cp310-win_arm64.whl", hash = "sha256:8fdf1f0951cc51e1b53903c14448d543cb595c701896f7698320f8db1bfab8ae", size = 94001, upload-time = "2025-07-26T01:10:54.315Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/54/26/141105c328ffcc6cb40b2dfb907e3c27c9f56f001aa52d335bb7a5c4aef1/lmdb-1.7.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:72e21d1752b7044fba00c77277e482683db72a62f0a2603c30fce66e3b26306d", size = 101023, upload-time = "2025-07-26T01:10:55.4Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/53/70/1208e98f9e917ef1041389b099809dbe3dfbe4660ce2d55ebc77bcfbe2a6/lmdb-1.7.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:2bea10371d95aaae0a48d0b34f76fd6580e73b20190e7834d83def04889b7e55", size = 98528, upload-time = "2025-07-26T01:10:56.476Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/45/87/0e31bc8ab34169415f4216d56d609595dce7ee88f8924712c4d6d1522e32/lmdb-1.7.3-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f9d7b6497be20e07b17f41a171b1663dfcfe5eac50fe003a256be5231ca5a121", size = 298041, upload-time = "2025-07-26T01:10:57.469Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f4/87/81d17a74498d2183b503e5e3f64892d8a5f94658ff271143cd732fe21bec/lmdb-1.7.3-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f939c5222f5725afb5dbdbbd9f2ef20cc1c9bad3696028bf9e5e62ba14e016d2", size = 299620, upload-time = "2025-07-26T01:10:58.944Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/11/e9/331a8aa028109a5275d7d3818a2696c44e26cd76d65863249bfc8adcfe10/lmdb-1.7.3-cp311-cp311-win_amd64.whl", hash = "sha256:2f42cc1d168fbd625e7981cf45fd8af15a2fb071aaa284e355c7036ad55e2ac7", size = 99276, upload-time = "2025-07-26T01:11:00.518Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a1/51/1e19c28d327bdca10d81df1e54546a190679a8816ab1cd5021e55c565997/lmdb-1.7.3-cp311-cp311-win_arm64.whl", hash = "sha256:8e1dc64e139bedaeb2aaa133d1139fdfad53e48a64cd772fb9a1d70fc403c1bc", size = 94120, upload-time = "2025-07-26T01:11:01.515Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/87/7c/3e90eeca144c430cc994a9b575215cd5a902327e248eede790ccc1731503/lmdb-1.7.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:a0c9d7b35adf001884687ca0c256f6c97aedfcc8bb118da92dcdef41e025cb92", size = 101186, upload-time = "2025-07-26T01:11:02.513Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/95/69/8e7542c81413f850ac0b7bc286fb50ea07f7fbed72c797788d8198c70531/lmdb-1.7.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:0522361cb78bc9e43003689f2dd3842e4ff8a9a8c11d41b942396f7b94602109", size = 98575, upload-time = "2025-07-26T01:11:03.661Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/99/c5/9e6d150f0070908e70ccae4dde5d2fc970a9de83b6ccd8214383bfb2b050/lmdb-1.7.3-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:92ee64c2de8185def1eac0d6272626cb57443895b5cd6f35eda848c0426e678e", size = 299883, upload-time = "2025-07-26T01:11:04.795Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/bc/02/98d3d0487aebfaa8fba1ced51bd73f86186293c698f36de585d65e6e8487/lmdb-1.7.3-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:50438f2b333f53539805d3565c0049507e743dff6280229980f1717ef5fcb0ac", size = 302994, upload-time = "2025-07-26T01:11:06.52Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b3/d7/e30633a116b8e3a887883e2ad7aea4231e67d47d9b1a6827a9a0d7d5aa45/lmdb-1.7.3-cp312-cp312-win_amd64.whl", hash = "sha256:a2cef8c30b6d5138ef2f7ed1fff661c497f988c5c7e38699295502bfd148a9ce", size = 99325, upload-time = "2025-07-26T01:11:07.881Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/10/5c/d417ac692eca2dcd5b2a8460389fa8883b5d7aafbccc0e48a206c6606365/lmdb-1.7.3-cp312-cp312-win_arm64.whl", hash = "sha256:4aa31d12b16fa0d4ab5d49a347a352eb65f8887ae52e9bfe1811b5daff83a064", size = 94184, upload-time = "2025-07-26T01:11:09.254Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/57/e8/2f89d1f2a1fd665dc663029dbeb19b9aa7e8fabaaa1491108112bcc44620/lmdb-1.7.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:3ef46b01b09afe8193449330ea6c0350891d1cdc43cbd20aec8ecdbe40def7d1", size = 101916, upload-time = "2025-07-26T01:11:10.432Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c5/57/e17901d6f0da20da6795d7fda35a68416b4b14e84873039c4b87870d45c6/lmdb-1.7.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:88d9b76a11e7458801064df8c837fe23ca28b714205dce05f17a58da3777db6a", size = 99340, upload-time = "2025-07-26T01:11:11.752Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b8/b7/7bfc65b5003ecc7f41fe6d4aeb3e425746cb2eda05d7314f712fca6d6094/lmdb-1.7.3-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:39dd69f3dc0910ed15ed5db708f195dc3c7f3275d1fb1a8f0a2a4b31db69cdd0", size = 301234, upload-time = "2025-07-26T01:11:12.951Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/68/65/b2e8aa0ddc984f3b8cf72223d878f0d7ff70f1744ad34b4917bc63b7dfe0/lmdb-1.7.3-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ad0dbec5586bb991c9ab05f5a26d950e8fa3583e632baf2604512f8059028139", size = 304101, upload-time = "2025-07-26T01:11:14.103Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d5/ea/059050689e8ad6c859d8bb7d88c7b4c183a7d54c6079aa6601e0935449ba/lmdb-1.7.3-cp313-cp313-win_amd64.whl", hash = "sha256:d25279380b28f01f8f1ef0f783488808bfc2b657c1705c6f383b77c9d04cd8a6", size = 99447, upload-time = "2025-07-26T01:11:15.606Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ad/38/b9e7e5807dbeaf66cd3326161f03ba2274c0774da15f213745708c30d924/lmdb-1.7.3-cp313-cp313-win_arm64.whl", hash = "sha256:1f22e79cc7cebc643c20f1699092bddb6d648d959f4f73becce0befb52690c97", size = 94146, upload-time = "2025-07-26T01:11:16.943Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d7/88/7931a5b49aa7f5d978ef8099efb5486270baf08f58c6b0df1ae274849c19/lmdb-1.7.3-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:825ba8e83b8b5bdfda362eac573a39f045c088c2a26dd3126c4e7fa63a04d331", size = 100817, upload-time = "2025-07-26T01:11:28.835Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4b/b2/b6d908fded5584a26a7f19b1fe6ad095e127daf843f9bcbff4004440b3be/lmdb-1.7.3-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:6cbeeb02e83f7b8551a646f77b87f5ac5b206399ed9f8f89d09e9ec56a139399", size = 98252, upload-time = "2025-07-26T01:11:29.846Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6b/21/5c6c8a3f7095d51373b6cc9094c436102c3903bac95b8d23960d374a97ed/lmdb-1.7.3-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f12b508807d353f877e437d762c923108f5233075b48abe2335dae48228153bf", size = 294774, upload-time = "2025-07-26T01:11:30.859Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c1/dd/14accf7b63b91cf281c19a4176b722b13565a7da7d79deba8acf43865a59/lmdb-1.7.3-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ce6955d1a387e19c3f793ff4a105e4dc56405a9b192644b418779fff1a3c0a8f", size = 296203, upload-time = "2025-07-26T01:11:32.335Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/28/a9/26ed714dc26774e367fd2c11520287d6958f8792bd44a14e91ae72ab86ea/lmdb-1.7.3-cp38-cp38-win_amd64.whl", hash = "sha256:75bfd980a6348b14337c5a63343d02dc57c3b9b6220103dd8e4ad4e9c43a2b9b", size = 99119, upload-time = "2025-07-26T01:11:33.779Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5e/91/15f5e6a59ab918376e1bd3df03cc7439366b097967f80094807f6016249f/lmdb-1.7.3-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:dace227f0d7aa22fd0bce4f43a549bee35b08e30c06285ae178dbf891b659d3b", size = 101035, upload-time = "2025-07-26T01:11:34.804Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7e/90/c203f59e8857156df4e88f381a752d278f62a6415b4ac72d5c41d82d04c7/lmdb-1.7.3-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:5e03d11cc29b33df4464a17955e42c73ec3bb2030b5005d7dbd7445c11545fc7", size = 98465, upload-time = "2025-07-26T01:11:36.112Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e7/3a/6e10d385e22d9d7f272fcac41fcddda1e309034e4206cd9ee42d81413637/lmdb-1.7.3-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bc669164f60e1508901085073b6deae24b147d16ed98bece4f7fa8b4bb9edfff", size = 294632, upload-time = "2025-07-26T01:11:37.764Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/25/5c/9997cdfc78f6e59a85db3519403710a3eeb55b4ca9aae37b48c615846e55/lmdb-1.7.3-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ea36951a87945ca8181a7393d1709f8d6ab8a187100485115cad1d8e03a41cae", size = 296066, upload-time = "2025-07-26T01:11:38.918Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8c/93/006220d20e13df4e87d01997888f64f7821b04729546c167944ee19f1241/lmdb-1.7.3-cp39-cp39-win_amd64.whl", hash = "sha256:020608dc83b81283c0ea9b2886f037ada113d7b2ee0f6cd5c3be57842e64d203", size = 99309, upload-time = "2025-07-26T01:11:40.17Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/12/cd/f2880a248d1d6eb2aec5fddd08956b473c4e01ebd793e277178c3a5960ad/lmdb-1.7.3-cp39-cp39-win_arm64.whl", hash = "sha256:413d55fe192ad221175f03941ee1bb71c2c059c615003b267a6ed203b713fabf", size = 94049, upload-time = "2025-07-26T01:11:41.182Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b9/ab/e6c2ba79107d20c7a45eae2d69e395691635e761b1e57f5a29c5f2fc42ff/lmdb-1.7.3-py3-none-any.whl", hash = "sha256:4cd7fb79b2c69dcf564d4eec36237e0efc303559176b60ff36fb47ccd3b8b5bb", size = 148538, upload-time = "2025-07-26T01:11:42.118Z" }, +] + +[[package]] +name = "markupsafe" +version = "2.1.5" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/87/5b/aae44c6655f3801e81aa3eef09dbbf012431987ba564d7231722f68df02d/MarkupSafe-2.1.5.tar.gz", hash = "sha256:d283d37a890ba4c1ae73ffadf8046435c76e7bc2247bbb63c00bd1a709c6544b", size = 19384, upload-time = "2024-02-02T16:31:22.863Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e4/54/ad5eb37bf9d51800010a74e4665425831a9db4e7c4e0fde4352e391e808e/MarkupSafe-2.1.5-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:a17a92de5231666cfbe003f0e4b9b3a7ae3afb1ec2845aadc2bacc93ff85febc", size = 18206, upload-time = "2024-02-02T16:30:04.105Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6a/4a/a4d49415e600bacae038c67f9fecc1d5433b9d3c71a4de6f33537b89654c/MarkupSafe-2.1.5-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:72b6be590cc35924b02c78ef34b467da4ba07e4e0f0454a2c5907f473fc50ce5", size = 14079, upload-time = "2024-02-02T16:30:06.5Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0a/7b/85681ae3c33c385b10ac0f8dd025c30af83c78cec1c37a6aa3b55e67f5ec/MarkupSafe-2.1.5-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e61659ba32cf2cf1481e575d0462554625196a1f2fc06a1c777d3f48e8865d46", size = 26620, upload-time = "2024-02-02T16:30:08.31Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7c/52/2b1b570f6b8b803cef5ac28fdf78c0da318916c7d2fe9402a84d591b394c/MarkupSafe-2.1.5-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2174c595a0d73a3080ca3257b40096db99799265e1c27cc5a610743acd86d62f", size = 25818, upload-time = "2024-02-02T16:30:09.577Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/29/fe/a36ba8c7ca55621620b2d7c585313efd10729e63ef81e4e61f52330da781/MarkupSafe-2.1.5-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ae2ad8ae6ebee9d2d94b17fb62763125f3f374c25618198f40cbb8b525411900", size = 25493, upload-time = "2024-02-02T16:30:11.488Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/60/ae/9c60231cdfda003434e8bd27282b1f4e197ad5a710c14bee8bea8a9ca4f0/MarkupSafe-2.1.5-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:075202fa5b72c86ad32dc7d0b56024ebdbcf2048c0ba09f1cde31bfdd57bcfff", size = 30630, upload-time = "2024-02-02T16:30:13.144Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/65/dc/1510be4d179869f5dafe071aecb3f1f41b45d37c02329dfba01ff59e5ac5/MarkupSafe-2.1.5-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:598e3276b64aff0e7b3451b72e94fa3c238d452e7ddcd893c3ab324717456bad", size = 29745, upload-time = "2024-02-02T16:30:14.222Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/30/39/8d845dd7d0b0613d86e0ef89549bfb5f61ed781f59af45fc96496e897f3a/MarkupSafe-2.1.5-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:fce659a462a1be54d2ffcacea5e3ba2d74daa74f30f5f143fe0c58636e355fdd", size = 30021, upload-time = "2024-02-02T16:30:16.032Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c7/5c/356a6f62e4f3c5fbf2602b4771376af22a3b16efa74eb8716fb4e328e01e/MarkupSafe-2.1.5-cp310-cp310-win32.whl", hash = "sha256:d9fad5155d72433c921b782e58892377c44bd6252b5af2f67f16b194987338a4", size = 16659, upload-time = "2024-02-02T16:30:17.079Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/69/48/acbf292615c65f0604a0c6fc402ce6d8c991276e16c80c46a8f758fbd30c/MarkupSafe-2.1.5-cp310-cp310-win_amd64.whl", hash = "sha256:bf50cd79a75d181c9181df03572cdce0fbb75cc353bc350712073108cba98de5", size = 17213, upload-time = "2024-02-02T16:30:18.251Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/11/e7/291e55127bb2ae67c64d66cef01432b5933859dfb7d6949daa721b89d0b3/MarkupSafe-2.1.5-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:629ddd2ca402ae6dbedfceeba9c46d5f7b2a61d9749597d4307f943ef198fc1f", size = 18219, upload-time = "2024-02-02T16:30:19.988Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6b/cb/aed7a284c00dfa7c0682d14df85ad4955a350a21d2e3b06d8240497359bf/MarkupSafe-2.1.5-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:5b7b716f97b52c5a14bffdf688f971b2d5ef4029127f1ad7a513973cfd818df2", size = 14098, upload-time = "2024-02-02T16:30:21.063Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1c/cf/35fe557e53709e93feb65575c93927942087e9b97213eabc3fe9d5b25a55/MarkupSafe-2.1.5-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6ec585f69cec0aa07d945b20805be741395e28ac1627333b1c5b0105962ffced", size = 29014, upload-time = "2024-02-02T16:30:22.926Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/97/18/c30da5e7a0e7f4603abfc6780574131221d9148f323752c2755d48abad30/MarkupSafe-2.1.5-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b91c037585eba9095565a3556f611e3cbfaa42ca1e865f7b8015fe5c7336d5a5", size = 28220, upload-time = "2024-02-02T16:30:24.76Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0c/40/2e73e7d532d030b1e41180807a80d564eda53babaf04d65e15c1cf897e40/MarkupSafe-2.1.5-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:7502934a33b54030eaf1194c21c692a534196063db72176b0c4028e140f8f32c", size = 27756, upload-time = "2024-02-02T16:30:25.877Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/18/46/5dca760547e8c59c5311b332f70605d24c99d1303dd9a6e1fc3ed0d73561/MarkupSafe-2.1.5-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:0e397ac966fdf721b2c528cf028494e86172b4feba51d65f81ffd65c63798f3f", size = 33988, upload-time = "2024-02-02T16:30:26.935Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6d/c5/27febe918ac36397919cd4a67d5579cbbfa8da027fa1238af6285bb368ea/MarkupSafe-2.1.5-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:c061bb86a71b42465156a3ee7bd58c8c2ceacdbeb95d05a99893e08b8467359a", size = 32718, upload-time = "2024-02-02T16:30:28.111Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f8/81/56e567126a2c2bc2684d6391332e357589a96a76cb9f8e5052d85cb0ead8/MarkupSafe-2.1.5-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:3a57fdd7ce31c7ff06cdfbf31dafa96cc533c21e443d57f5b1ecc6cdc668ec7f", size = 33317, upload-time = "2024-02-02T16:30:29.214Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/00/0b/23f4b2470accb53285c613a3ab9ec19dc944eaf53592cb6d9e2af8aa24cc/MarkupSafe-2.1.5-cp311-cp311-win32.whl", hash = "sha256:397081c1a0bfb5124355710fe79478cdbeb39626492b15d399526ae53422b906", size = 16670, upload-time = "2024-02-02T16:30:30.915Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b7/a2/c78a06a9ec6d04b3445a949615c4c7ed86a0b2eb68e44e7541b9d57067cc/MarkupSafe-2.1.5-cp311-cp311-win_amd64.whl", hash = "sha256:2b7c57a4dfc4f16f7142221afe5ba4e093e09e728ca65c51f5620c9aaeb9a617", size = 17224, upload-time = "2024-02-02T16:30:32.09Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/53/bd/583bf3e4c8d6a321938c13f49d44024dbe5ed63e0a7ba127e454a66da974/MarkupSafe-2.1.5-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:8dec4936e9c3100156f8a2dc89c4b88d5c435175ff03413b443469c7c8c5f4d1", size = 18215, upload-time = "2024-02-02T16:30:33.081Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/48/d6/e7cd795fc710292c3af3a06d80868ce4b02bfbbf370b7cee11d282815a2a/MarkupSafe-2.1.5-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:3c6b973f22eb18a789b1460b4b91bf04ae3f0c4234a0a6aa6b0a92f6f7b951d4", size = 14069, upload-time = "2024-02-02T16:30:34.148Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/51/b5/5d8ec796e2a08fc814a2c7d2584b55f889a55cf17dd1a90f2beb70744e5c/MarkupSafe-2.1.5-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ac07bad82163452a6884fe8fa0963fb98c2346ba78d779ec06bd7a6262132aee", size = 29452, upload-time = "2024-02-02T16:30:35.149Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0a/0d/2454f072fae3b5a137c119abf15465d1771319dfe9e4acbb31722a0fff91/MarkupSafe-2.1.5-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f5dfb42c4604dddc8e4305050aa6deb084540643ed5804d7455b5df8fe16f5e5", size = 28462, upload-time = "2024-02-02T16:30:36.166Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2d/75/fd6cb2e68780f72d47e6671840ca517bda5ef663d30ada7616b0462ad1e3/MarkupSafe-2.1.5-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ea3d8a3d18833cf4304cd2fc9cbb1efe188ca9b5efef2bdac7adc20594a0e46b", size = 27869, upload-time = "2024-02-02T16:30:37.834Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b0/81/147c477391c2750e8fc7705829f7351cf1cd3be64406edcf900dc633feb2/MarkupSafe-2.1.5-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:d050b3361367a06d752db6ead6e7edeb0009be66bc3bae0ee9d97fb326badc2a", size = 33906, upload-time = "2024-02-02T16:30:39.366Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8b/ff/9a52b71839d7a256b563e85d11050e307121000dcebc97df120176b3ad93/MarkupSafe-2.1.5-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:bec0a414d016ac1a18862a519e54b2fd0fc8bbfd6890376898a6c0891dd82e9f", size = 32296, upload-time = "2024-02-02T16:30:40.413Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/88/07/2dc76aa51b481eb96a4c3198894f38b480490e834479611a4053fbf08623/MarkupSafe-2.1.5-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:58c98fee265677f63a4385256a6d7683ab1832f3ddd1e66fe948d5880c21a169", size = 33038, upload-time = "2024-02-02T16:30:42.243Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/96/0c/620c1fb3661858c0e37eb3cbffd8c6f732a67cd97296f725789679801b31/MarkupSafe-2.1.5-cp312-cp312-win32.whl", hash = "sha256:8590b4ae07a35970728874632fed7bd57b26b0102df2d2b233b6d9d82f6c62ad", size = 16572, upload-time = "2024-02-02T16:30:43.326Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3f/14/c3554d512d5f9100a95e737502f4a2323a1959f6d0d01e0d0997b35f7b10/MarkupSafe-2.1.5-cp312-cp312-win_amd64.whl", hash = "sha256:823b65d8706e32ad2df51ed89496147a42a2a6e01c13cfb6ffb8b1e92bc910bb", size = 17127, upload-time = "2024-02-02T16:30:44.418Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f8/ff/2c942a82c35a49df5de3a630ce0a8456ac2969691b230e530ac12314364c/MarkupSafe-2.1.5-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:656f7526c69fac7f600bd1f400991cc282b417d17539a1b228617081106feb4a", size = 18192, upload-time = "2024-02-02T16:30:57.715Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4f/14/6f294b9c4f969d0c801a4615e221c1e084722ea6114ab2114189c5b8cbe0/MarkupSafe-2.1.5-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:97cafb1f3cbcd3fd2b6fbfb99ae11cdb14deea0736fc2b0952ee177f2b813a46", size = 14072, upload-time = "2024-02-02T16:30:58.844Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/81/d4/fd74714ed30a1dedd0b82427c02fa4deec64f173831ec716da11c51a50aa/MarkupSafe-2.1.5-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1f3fbcb7ef1f16e48246f704ab79d79da8a46891e2da03f8783a5b6fa41a9532", size = 26928, upload-time = "2024-02-02T16:30:59.922Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c7/bd/50319665ce81bb10e90d1cf76f9e1aa269ea6f7fa30ab4521f14d122a3df/MarkupSafe-2.1.5-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fa9db3f79de01457b03d4f01b34cf91bc0048eb2c3846ff26f66687c2f6d16ab", size = 26106, upload-time = "2024-02-02T16:31:01.582Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4c/6f/f2b0f675635b05f6afd5ea03c094557bdb8622fa8e673387444fe8d8e787/MarkupSafe-2.1.5-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ffee1f21e5ef0d712f9033568f8344d5da8cc2869dbd08d87c84656e6a2d2f68", size = 25781, upload-time = "2024-02-02T16:31:02.71Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/51/e0/393467cf899b34a9d3678e78961c2c8cdf49fb902a959ba54ece01273fb1/MarkupSafe-2.1.5-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:5dedb4db619ba5a2787a94d877bc8ffc0566f92a01c0ef214865e54ecc9ee5e0", size = 30518, upload-time = "2024-02-02T16:31:04.392Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f6/02/5437e2ad33047290dafced9df741d9efc3e716b75583bbd73a9984f1b6f7/MarkupSafe-2.1.5-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:30b600cf0a7ac9234b2638fbc0fb6158ba5bdcdf46aeb631ead21248b9affbc4", size = 29669, upload-time = "2024-02-02T16:31:05.53Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0e/7d/968284145ffd9d726183ed6237c77938c021abacde4e073020f920e060b2/MarkupSafe-2.1.5-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:8dd717634f5a044f860435c1d8c16a270ddf0ef8588d4887037c5028b859b0c3", size = 29933, upload-time = "2024-02-02T16:31:06.636Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/bf/f3/ecb00fc8ab02b7beae8699f34db9357ae49d9f21d4d3de6f305f34fa949e/MarkupSafe-2.1.5-cp38-cp38-win32.whl", hash = "sha256:daa4ee5a243f0f20d528d939d06670a298dd39b1ad5f8a72a4275124a7819eff", size = 16656, upload-time = "2024-02-02T16:31:07.767Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/92/21/357205f03514a49b293e214ac39de01fadd0970a6e05e4bf1ddd0ffd0881/MarkupSafe-2.1.5-cp38-cp38-win_amd64.whl", hash = "sha256:619bc166c4f2de5caa5a633b8b7326fbe98e0ccbfacabd87268a2b15ff73a029", size = 17206, upload-time = "2024-02-02T16:31:08.843Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0f/31/780bb297db036ba7b7bbede5e1d7f1e14d704ad4beb3ce53fb495d22bc62/MarkupSafe-2.1.5-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:7a68b554d356a91cce1236aa7682dc01df0edba8d043fd1ce607c49dd3c1edcf", size = 18193, upload-time = "2024-02-02T16:31:10.155Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6c/77/d77701bbef72892affe060cdacb7a2ed7fd68dae3b477a8642f15ad3b132/MarkupSafe-2.1.5-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:db0b55e0f3cc0be60c1f19efdde9a637c32740486004f20d1cff53c3c0ece4d2", size = 14073, upload-time = "2024-02-02T16:31:11.442Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d9/a7/1e558b4f78454c8a3a0199292d96159eb4d091f983bc35ef258314fe7269/MarkupSafe-2.1.5-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3e53af139f8579a6d5f7b76549125f0d94d7e630761a2111bc431fd820e163b8", size = 26486, upload-time = "2024-02-02T16:31:12.488Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5f/5a/360da85076688755ea0cceb92472923086993e86b5613bbae9fbc14136b0/MarkupSafe-2.1.5-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:17b950fccb810b3293638215058e432159d2b71005c74371d784862b7e4683f3", size = 25685, upload-time = "2024-02-02T16:31:13.726Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6a/18/ae5a258e3401f9b8312f92b028c54d7026a97ec3ab20bfaddbdfa7d8cce8/MarkupSafe-2.1.5-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:4c31f53cdae6ecfa91a77820e8b151dba54ab528ba65dfd235c80b086d68a465", size = 25338, upload-time = "2024-02-02T16:31:14.812Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0b/cc/48206bd61c5b9d0129f4d75243b156929b04c94c09041321456fd06a876d/MarkupSafe-2.1.5-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:bff1b4290a66b490a2f4719358c0cdcd9bafb6b8f061e45c7a2460866bf50c2e", size = 30439, upload-time = "2024-02-02T16:31:15.946Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d1/06/a41c112ab9ffdeeb5f77bc3e331fdadf97fa65e52e44ba31880f4e7f983c/MarkupSafe-2.1.5-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:bc1667f8b83f48511b94671e0e441401371dfd0f0a795c7daa4a3cd1dde55bea", size = 29531, upload-time = "2024-02-02T16:31:17.13Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/02/8c/ab9a463301a50dab04d5472e998acbd4080597abc048166ded5c7aa768c8/MarkupSafe-2.1.5-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:5049256f536511ee3f7e1b3f87d1d1209d327e818e6ae1365e8653d7e3abb6a6", size = 29823, upload-time = "2024-02-02T16:31:18.247Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/bc/29/9bc18da763496b055d8e98ce476c8e718dcfd78157e17f555ce6dd7d0895/MarkupSafe-2.1.5-cp39-cp39-win32.whl", hash = "sha256:00e046b6dd71aa03a41079792f8473dc494d564611a8f89bbbd7cb93295ebdcf", size = 16658, upload-time = "2024-02-02T16:31:19.583Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f6/f8/4da07de16f10551ca1f640c92b5f316f9394088b183c6a57183df6de5ae4/MarkupSafe-2.1.5-cp39-cp39-win_amd64.whl", hash = "sha256:fa173ec60341d6bb97a89f5ea19c85c5643c1e7dedebc22f5181eb73573142c5", size = 17211, upload-time = "2024-02-02T16:31:20.96Z" }, +] + +[[package]] +name = "markupsafe" +version = "3.0.2" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.13.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'darwin')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'darwin')", + "(python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform == 'darwin')", + "(python_full_version == '3.10.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'darwin')", + "(python_full_version == '3.9.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform == 'darwin')", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b2/97/5d42485e71dfc078108a86d6de8fa46db44a1a9295e89c5d6d4a06e23a62/markupsafe-3.0.2.tar.gz", hash = "sha256:ee55d3edf80167e48ea11a923c7386f4669df67d7994554387f84e7d8b0a2bf0", size = 20537, upload-time = "2024-10-18T15:21:54.129Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/04/90/d08277ce111dd22f77149fd1a5d4653eeb3b3eaacbdfcbae5afb2600eebd/MarkupSafe-3.0.2-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:7e94c425039cde14257288fd61dcfb01963e658efbc0ff54f5306b06054700f8", size = 14357, upload-time = "2024-10-18T15:20:51.44Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/04/e1/6e2194baeae0bca1fae6629dc0cbbb968d4d941469cbab11a3872edff374/MarkupSafe-3.0.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:9e2d922824181480953426608b81967de705c3cef4d1af983af849d7bd619158", size = 12393, upload-time = "2024-10-18T15:20:52.426Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1d/69/35fa85a8ece0a437493dc61ce0bb6d459dcba482c34197e3efc829aa357f/MarkupSafe-3.0.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:38a9ef736c01fccdd6600705b09dc574584b89bea478200c5fbf112a6b0d5579", size = 21732, upload-time = "2024-10-18T15:20:53.578Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/22/35/137da042dfb4720b638d2937c38a9c2df83fe32d20e8c8f3185dbfef05f7/MarkupSafe-3.0.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bbcb445fa71794da8f178f0f6d66789a28d7319071af7a496d4d507ed566270d", size = 20866, upload-time = "2024-10-18T15:20:55.06Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/29/28/6d029a903727a1b62edb51863232152fd335d602def598dade38996887f0/MarkupSafe-3.0.2-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:57cb5a3cf367aeb1d316576250f65edec5bb3be939e9247ae594b4bcbc317dfb", size = 20964, upload-time = "2024-10-18T15:20:55.906Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/cc/cd/07438f95f83e8bc028279909d9c9bd39e24149b0d60053a97b2bc4f8aa51/MarkupSafe-3.0.2-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:3809ede931876f5b2ec92eef964286840ed3540dadf803dd570c3b7e13141a3b", size = 21977, upload-time = "2024-10-18T15:20:57.189Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/29/01/84b57395b4cc062f9c4c55ce0df7d3108ca32397299d9df00fedd9117d3d/MarkupSafe-3.0.2-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:e07c3764494e3776c602c1e78e298937c3315ccc9043ead7e685b7f2b8d47b3c", size = 21366, upload-time = "2024-10-18T15:20:58.235Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/bd/6e/61ebf08d8940553afff20d1fb1ba7294b6f8d279df9fd0c0db911b4bbcfd/MarkupSafe-3.0.2-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:b424c77b206d63d500bcb69fa55ed8d0e6a3774056bdc4839fc9298a7edca171", size = 21091, upload-time = "2024-10-18T15:20:59.235Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/11/23/ffbf53694e8c94ebd1e7e491de185124277964344733c45481f32ede2499/MarkupSafe-3.0.2-cp310-cp310-win32.whl", hash = "sha256:fcabf5ff6eea076f859677f5f0b6b5c1a51e70a376b0579e0eadef8db48c6b50", size = 15065, upload-time = "2024-10-18T15:21:00.307Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/44/06/e7175d06dd6e9172d4a69a72592cb3f7a996a9c396eee29082826449bbc3/MarkupSafe-3.0.2-cp310-cp310-win_amd64.whl", hash = "sha256:6af100e168aa82a50e186c82875a5893c5597a0c1ccdb0d8b40240b1f28b969a", size = 15514, upload-time = "2024-10-18T15:21:01.122Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6b/28/bbf83e3f76936960b850435576dd5e67034e200469571be53f69174a2dfd/MarkupSafe-3.0.2-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:9025b4018f3a1314059769c7bf15441064b2207cb3f065e6ea1e7359cb46db9d", size = 14353, upload-time = "2024-10-18T15:21:02.187Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6c/30/316d194b093cde57d448a4c3209f22e3046c5bb2fb0820b118292b334be7/MarkupSafe-3.0.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:93335ca3812df2f366e80509ae119189886b0f3c2b81325d39efdb84a1e2ae93", size = 12392, upload-time = "2024-10-18T15:21:02.941Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f2/96/9cdafba8445d3a53cae530aaf83c38ec64c4d5427d975c974084af5bc5d2/MarkupSafe-3.0.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2cb8438c3cbb25e220c2ab33bb226559e7afb3baec11c4f218ffa7308603c832", size = 23984, upload-time = "2024-10-18T15:21:03.953Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f1/a4/aefb044a2cd8d7334c8a47d3fb2c9f328ac48cb349468cc31c20b539305f/MarkupSafe-3.0.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a123e330ef0853c6e822384873bef7507557d8e4a082961e1defa947aa59ba84", size = 23120, upload-time = "2024-10-18T15:21:06.495Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8d/21/5e4851379f88f3fad1de30361db501300d4f07bcad047d3cb0449fc51f8c/MarkupSafe-3.0.2-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1e084f686b92e5b83186b07e8a17fc09e38fff551f3602b249881fec658d3eca", size = 23032, upload-time = "2024-10-18T15:21:07.295Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/00/7b/e92c64e079b2d0d7ddf69899c98842f3f9a60a1ae72657c89ce2655c999d/MarkupSafe-3.0.2-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:d8213e09c917a951de9d09ecee036d5c7d36cb6cb7dbaece4c71a60d79fb9798", size = 24057, upload-time = "2024-10-18T15:21:08.073Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f9/ac/46f960ca323037caa0a10662ef97d0a4728e890334fc156b9f9e52bcc4ca/MarkupSafe-3.0.2-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:5b02fb34468b6aaa40dfc198d813a641e3a63b98c2b05a16b9f80b7ec314185e", size = 23359, upload-time = "2024-10-18T15:21:09.318Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/69/84/83439e16197337b8b14b6a5b9c2105fff81d42c2a7c5b58ac7b62ee2c3b1/MarkupSafe-3.0.2-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:0bff5e0ae4ef2e1ae4fdf2dfd5b76c75e5c2fa4132d05fc1b0dabcd20c7e28c4", size = 23306, upload-time = "2024-10-18T15:21:10.185Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/9a/34/a15aa69f01e2181ed8d2b685c0d2f6655d5cca2c4db0ddea775e631918cd/MarkupSafe-3.0.2-cp311-cp311-win32.whl", hash = "sha256:6c89876f41da747c8d3677a2b540fb32ef5715f97b66eeb0c6b66f5e3ef6f59d", size = 15094, upload-time = "2024-10-18T15:21:11.005Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/da/b8/3a3bd761922d416f3dc5d00bfbed11f66b1ab89a0c2b6e887240a30b0f6b/MarkupSafe-3.0.2-cp311-cp311-win_amd64.whl", hash = "sha256:70a87b411535ccad5ef2f1df5136506a10775d267e197e4cf531ced10537bd6b", size = 15521, upload-time = "2024-10-18T15:21:12.911Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/22/09/d1f21434c97fc42f09d290cbb6350d44eb12f09cc62c9476effdb33a18aa/MarkupSafe-3.0.2-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:9778bd8ab0a994ebf6f84c2b949e65736d5575320a17ae8984a77fab08db94cf", size = 14274, upload-time = "2024-10-18T15:21:13.777Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6b/b0/18f76bba336fa5aecf79d45dcd6c806c280ec44538b3c13671d49099fdd0/MarkupSafe-3.0.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:846ade7b71e3536c4e56b386c2a47adf5741d2d8b94ec9dc3e92e5e1ee1e2225", size = 12348, upload-time = "2024-10-18T15:21:14.822Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e0/25/dd5c0f6ac1311e9b40f4af06c78efde0f3b5cbf02502f8ef9501294c425b/MarkupSafe-3.0.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1c99d261bd2d5f6b59325c92c73df481e05e57f19837bdca8413b9eac4bd8028", size = 24149, upload-time = "2024-10-18T15:21:15.642Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f3/f0/89e7aadfb3749d0f52234a0c8c7867877876e0a20b60e2188e9850794c17/MarkupSafe-3.0.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e17c96c14e19278594aa4841ec148115f9c7615a47382ecb6b82bd8fea3ab0c8", size = 23118, upload-time = "2024-10-18T15:21:17.133Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d5/da/f2eeb64c723f5e3777bc081da884b414671982008c47dcc1873d81f625b6/MarkupSafe-3.0.2-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:88416bd1e65dcea10bc7569faacb2c20ce071dd1f87539ca2ab364bf6231393c", size = 22993, upload-time = "2024-10-18T15:21:18.064Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/da/0e/1f32af846df486dce7c227fe0f2398dc7e2e51d4a370508281f3c1c5cddc/MarkupSafe-3.0.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:2181e67807fc2fa785d0592dc2d6206c019b9502410671cc905d132a92866557", size = 24178, upload-time = "2024-10-18T15:21:18.859Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c4/f6/bb3ca0532de8086cbff5f06d137064c8410d10779c4c127e0e47d17c0b71/MarkupSafe-3.0.2-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:52305740fe773d09cffb16f8ed0427942901f00adedac82ec8b67752f58a1b22", size = 23319, upload-time = "2024-10-18T15:21:19.671Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a2/82/8be4c96ffee03c5b4a034e60a31294daf481e12c7c43ab8e34a1453ee48b/MarkupSafe-3.0.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:ad10d3ded218f1039f11a75f8091880239651b52e9bb592ca27de44eed242a48", size = 23352, upload-time = "2024-10-18T15:21:20.971Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/51/ae/97827349d3fcffee7e184bdf7f41cd6b88d9919c80f0263ba7acd1bbcb18/MarkupSafe-3.0.2-cp312-cp312-win32.whl", hash = "sha256:0f4ca02bea9a23221c0182836703cbf8930c5e9454bacce27e767509fa286a30", size = 15097, upload-time = "2024-10-18T15:21:22.646Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c1/80/a61f99dc3a936413c3ee4e1eecac96c0da5ed07ad56fd975f1a9da5bc630/MarkupSafe-3.0.2-cp312-cp312-win_amd64.whl", hash = "sha256:8e06879fc22a25ca47312fbe7c8264eb0b662f6db27cb2d3bbbc74b1df4b9b87", size = 15601, upload-time = "2024-10-18T15:21:23.499Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/83/0e/67eb10a7ecc77a0c2bbe2b0235765b98d164d81600746914bebada795e97/MarkupSafe-3.0.2-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:ba9527cdd4c926ed0760bc301f6728ef34d841f405abf9d4f959c478421e4efd", size = 14274, upload-time = "2024-10-18T15:21:24.577Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2b/6d/9409f3684d3335375d04e5f05744dfe7e9f120062c9857df4ab490a1031a/MarkupSafe-3.0.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:f8b3d067f2e40fe93e1ccdd6b2e1d16c43140e76f02fb1319a05cf2b79d99430", size = 12352, upload-time = "2024-10-18T15:21:25.382Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d2/f5/6eadfcd3885ea85fe2a7c128315cc1bb7241e1987443d78c8fe712d03091/MarkupSafe-3.0.2-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:569511d3b58c8791ab4c2e1285575265991e6d8f8700c7be0e88f86cb0672094", size = 24122, upload-time = "2024-10-18T15:21:26.199Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0c/91/96cf928db8236f1bfab6ce15ad070dfdd02ed88261c2afafd4b43575e9e9/MarkupSafe-3.0.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:15ab75ef81add55874e7ab7055e9c397312385bd9ced94920f2802310c930396", size = 23085, upload-time = "2024-10-18T15:21:27.029Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c2/cf/c9d56af24d56ea04daae7ac0940232d31d5a8354f2b457c6d856b2057d69/MarkupSafe-3.0.2-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f3818cb119498c0678015754eba762e0d61e5b52d34c8b13d770f0719f7b1d79", size = 22978, upload-time = "2024-10-18T15:21:27.846Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2a/9f/8619835cd6a711d6272d62abb78c033bda638fdc54c4e7f4272cf1c0962b/MarkupSafe-3.0.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:cdb82a876c47801bb54a690c5ae105a46b392ac6099881cdfb9f6e95e4014c6a", size = 24208, upload-time = "2024-10-18T15:21:28.744Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f9/bf/176950a1792b2cd2102b8ffeb5133e1ed984547b75db47c25a67d3359f77/MarkupSafe-3.0.2-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:cabc348d87e913db6ab4aa100f01b08f481097838bdddf7c7a84b7575b7309ca", size = 23357, upload-time = "2024-10-18T15:21:29.545Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ce/4f/9a02c1d335caabe5c4efb90e1b6e8ee944aa245c1aaaab8e8a618987d816/MarkupSafe-3.0.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:444dcda765c8a838eaae23112db52f1efaf750daddb2d9ca300bcae1039adc5c", size = 23344, upload-time = "2024-10-18T15:21:30.366Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ee/55/c271b57db36f748f0e04a759ace9f8f759ccf22b4960c270c78a394f58be/MarkupSafe-3.0.2-cp313-cp313-win32.whl", hash = "sha256:bcf3e58998965654fdaff38e58584d8937aa3096ab5354d493c77d1fdd66d7a1", size = 15101, upload-time = "2024-10-18T15:21:31.207Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/29/88/07df22d2dd4df40aba9f3e402e6dc1b8ee86297dddbad4872bd5e7b0094f/MarkupSafe-3.0.2-cp313-cp313-win_amd64.whl", hash = "sha256:e6a2a455bd412959b57a172ce6328d2dd1f01cb2135efda2e4576e8a23fa3b0f", size = 15603, upload-time = "2024-10-18T15:21:32.032Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/62/6a/8b89d24db2d32d433dffcd6a8779159da109842434f1dd2f6e71f32f738c/MarkupSafe-3.0.2-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:b5a6b3ada725cea8a5e634536b1b01c30bcdcd7f9c6fff4151548d5bf6b3a36c", size = 14510, upload-time = "2024-10-18T15:21:33.625Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7a/06/a10f955f70a2e5a9bf78d11a161029d278eeacbd35ef806c3fd17b13060d/MarkupSafe-3.0.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:a904af0a6162c73e3edcb969eeeb53a63ceeb5d8cf642fade7d39e7963a22ddb", size = 12486, upload-time = "2024-10-18T15:21:34.611Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/34/cf/65d4a571869a1a9078198ca28f39fba5fbb910f952f9dbc5220afff9f5e6/MarkupSafe-3.0.2-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4aa4e5faecf353ed117801a068ebab7b7e09ffb6e1d5e412dc852e0da018126c", size = 25480, upload-time = "2024-10-18T15:21:35.398Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0c/e3/90e9651924c430b885468b56b3d597cabf6d72be4b24a0acd1fa0e12af67/MarkupSafe-3.0.2-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c0ef13eaeee5b615fb07c9a7dadb38eac06a0608b41570d8ade51c56539e509d", size = 23914, upload-time = "2024-10-18T15:21:36.231Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/66/8c/6c7cf61f95d63bb866db39085150df1f2a5bd3335298f14a66b48e92659c/MarkupSafe-3.0.2-cp313-cp313t-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d16a81a06776313e817c951135cf7340a3e91e8c1ff2fac444cfd75fffa04afe", size = 23796, upload-time = "2024-10-18T15:21:37.073Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/bb/35/cbe9238ec3f47ac9a7c8b3df7a808e7cb50fe149dc7039f5f454b3fba218/MarkupSafe-3.0.2-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:6381026f158fdb7c72a168278597a5e3a5222e83ea18f543112b2662a9b699c5", size = 25473, upload-time = "2024-10-18T15:21:37.932Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e6/32/7621a4382488aa283cc05e8984a9c219abad3bca087be9ec77e89939ded9/MarkupSafe-3.0.2-cp313-cp313t-musllinux_1_2_i686.whl", hash = "sha256:3d79d162e7be8f996986c064d1c7c817f6df3a77fe3d6859f6f9e7be4b8c213a", size = 24114, upload-time = "2024-10-18T15:21:39.799Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0d/80/0985960e4b89922cb5a0bac0ed39c5b96cbc1a536a99f30e8c220a996ed9/MarkupSafe-3.0.2-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:131a3c7689c85f5ad20f9f6fb1b866f402c445b220c19fe4308c0b147ccd2ad9", size = 24098, upload-time = "2024-10-18T15:21:40.813Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/82/78/fedb03c7d5380df2427038ec8d973587e90561b2d90cd472ce9254cf348b/MarkupSafe-3.0.2-cp313-cp313t-win32.whl", hash = "sha256:ba8062ed2cf21c07a9e295d5b8a2a5ce678b913b45fdf68c32d95d6c1291e0b6", size = 15208, upload-time = "2024-10-18T15:21:41.814Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4f/65/6079a46068dfceaeabb5dcad6d674f5f5c61a6fa5673746f42a9f4c233b3/MarkupSafe-3.0.2-cp313-cp313t-win_amd64.whl", hash = "sha256:e444a31f8db13eb18ada366ab3cf45fd4b31e4db1236a4448f68778c1d1a5a2f", size = 15739, upload-time = "2024-10-18T15:21:42.784Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a7/ea/9b1530c3fdeeca613faeb0fb5cbcf2389d816072fab72a71b45749ef6062/MarkupSafe-3.0.2-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:eaa0a10b7f72326f1372a713e73c3f739b524b3af41feb43e4921cb529f5929a", size = 14344, upload-time = "2024-10-18T15:21:43.721Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4b/c2/fbdbfe48848e7112ab05e627e718e854d20192b674952d9042ebd8c9e5de/MarkupSafe-3.0.2-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:48032821bbdf20f5799ff537c7ac3d1fba0ba032cfc06194faffa8cda8b560ff", size = 12389, upload-time = "2024-10-18T15:21:44.666Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f0/25/7a7c6e4dbd4f867d95d94ca15449e91e52856f6ed1905d58ef1de5e211d0/MarkupSafe-3.0.2-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1a9d3f5f0901fdec14d8d2f66ef7d035f2157240a433441719ac9a3fba440b13", size = 21607, upload-time = "2024-10-18T15:21:45.452Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/53/8f/f339c98a178f3c1e545622206b40986a4c3307fe39f70ccd3d9df9a9e425/MarkupSafe-3.0.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:88b49a3b9ff31e19998750c38e030fc7bb937398b1f78cfa599aaef92d693144", size = 20728, upload-time = "2024-10-18T15:21:46.295Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1a/03/8496a1a78308456dbd50b23a385c69b41f2e9661c67ea1329849a598a8f9/MarkupSafe-3.0.2-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:cfad01eed2c2e0c01fd0ecd2ef42c492f7f93902e39a42fc9ee1692961443a29", size = 20826, upload-time = "2024-10-18T15:21:47.134Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e6/cf/0a490a4bd363048c3022f2f475c8c05582179bb179defcee4766fb3dcc18/MarkupSafe-3.0.2-cp39-cp39-musllinux_1_2_aarch64.whl", hash = "sha256:1225beacc926f536dc82e45f8a4d68502949dc67eea90eab715dea3a21c1b5f0", size = 21843, upload-time = "2024-10-18T15:21:48.334Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/19/a3/34187a78613920dfd3cdf68ef6ce5e99c4f3417f035694074beb8848cd77/MarkupSafe-3.0.2-cp39-cp39-musllinux_1_2_i686.whl", hash = "sha256:3169b1eefae027567d1ce6ee7cae382c57fe26e82775f460f0b2778beaad66c0", size = 21219, upload-time = "2024-10-18T15:21:49.587Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/17/d8/5811082f85bb88410ad7e452263af048d685669bbbfb7b595e8689152498/MarkupSafe-3.0.2-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:eb7972a85c54febfb25b5c4b4f3af4dcc731994c7da0d8a0b4a6eb0640e1d178", size = 20946, upload-time = "2024-10-18T15:21:50.441Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7c/31/bd635fb5989440d9365c5e3c47556cfea121c7803f5034ac843e8f37c2f2/MarkupSafe-3.0.2-cp39-cp39-win32.whl", hash = "sha256:8c4e8c3ce11e1f92f6536ff07154f9d49677ebaaafc32db9db4620bc11ed480f", size = 15063, upload-time = "2024-10-18T15:21:51.385Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b3/73/085399401383ce949f727afec55ec3abd76648d04b9f22e1c0e99cb4bec3/MarkupSafe-3.0.2-cp39-cp39-win_amd64.whl", hash = "sha256:6e296a513ca3d94054c2c881cc913116e90fd030ad1c656b3869762b754f5f8a", size = 15506, upload-time = "2024-10-18T15:21:52.974Z" }, +] + +[[package]] +name = "mpmath" +version = "1.3.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106, upload-time = "2023-03-07T16:47:11.061Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, +] + +[[package]] +name = "msgpack" +version = "1.1.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/45/b1/ea4f68038a18c77c9467400d166d74c4ffa536f34761f7983a104357e614/msgpack-1.1.1.tar.gz", hash = "sha256:77b79ce34a2bdab2594f490c8e80dd62a02d650b91a75159a63ec413b8d104cd", size = 173555, upload-time = "2025-06-13T06:52:51.324Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/33/52/f30da112c1dc92cf64f57d08a273ac771e7b29dea10b4b30369b2d7e8546/msgpack-1.1.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:353b6fc0c36fde68b661a12949d7d49f8f51ff5fa019c1e47c87c4ff34b080ed", size = 81799, upload-time = "2025-06-13T06:51:37.228Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e4/35/7bfc0def2f04ab4145f7f108e3563f9b4abae4ab0ed78a61f350518cc4d2/msgpack-1.1.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:79c408fcf76a958491b4e3b103d1c417044544b68e96d06432a189b43d1215c8", size = 78278, upload-time = "2025-06-13T06:51:38.534Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e8/c5/df5d6c1c39856bc55f800bf82778fd4c11370667f9b9e9d51b2f5da88f20/msgpack-1.1.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:78426096939c2c7482bf31ef15ca219a9e24460289c00dd0b94411040bb73ad2", size = 402805, upload-time = "2025-06-13T06:51:39.538Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/20/8e/0bb8c977efecfe6ea7116e2ed73a78a8d32a947f94d272586cf02a9757db/msgpack-1.1.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8b17ba27727a36cb73aabacaa44b13090feb88a01d012c0f4be70c00f75048b4", size = 408642, upload-time = "2025-06-13T06:51:41.092Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/59/a1/731d52c1aeec52006be6d1f8027c49fdc2cfc3ab7cbe7c28335b2910d7b6/msgpack-1.1.1-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:7a17ac1ea6ec3c7687d70201cfda3b1e8061466f28f686c24f627cae4ea8efd0", size = 395143, upload-time = "2025-06-13T06:51:42.575Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2b/92/b42911c52cda2ba67a6418ffa7d08969edf2e760b09015593c8a8a27a97d/msgpack-1.1.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:88d1e966c9235c1d4e2afac21ca83933ba59537e2e2727a999bf3f515ca2af26", size = 395986, upload-time = "2025-06-13T06:51:43.807Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/61/dc/8ae165337e70118d4dab651b8b562dd5066dd1e6dd57b038f32ebc3e2f07/msgpack-1.1.1-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:f6d58656842e1b2ddbe07f43f56b10a60f2ba5826164910968f5933e5178af75", size = 402682, upload-time = "2025-06-13T06:51:45.534Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/58/27/555851cb98dcbd6ce041df1eacb25ac30646575e9cd125681aa2f4b1b6f1/msgpack-1.1.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:96decdfc4adcbc087f5ea7ebdcfd3dee9a13358cae6e81d54be962efc38f6338", size = 406368, upload-time = "2025-06-13T06:51:46.97Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d4/64/39a26add4ce16f24e99eabb9005e44c663db00e3fce17d4ae1ae9d61df99/msgpack-1.1.1-cp310-cp310-win32.whl", hash = "sha256:6640fd979ca9a212e4bcdf6eb74051ade2c690b862b679bfcb60ae46e6dc4bfd", size = 65004, upload-time = "2025-06-13T06:51:48.582Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7d/18/73dfa3e9d5d7450d39debde5b0d848139f7de23bd637a4506e36c9800fd6/msgpack-1.1.1-cp310-cp310-win_amd64.whl", hash = "sha256:8b65b53204fe1bd037c40c4148d00ef918eb2108d24c9aaa20bc31f9810ce0a8", size = 71548, upload-time = "2025-06-13T06:51:49.558Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7f/83/97f24bf9848af23fe2ba04380388216defc49a8af6da0c28cc636d722502/msgpack-1.1.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:71ef05c1726884e44f8b1d1773604ab5d4d17729d8491403a705e649116c9558", size = 82728, upload-time = "2025-06-13T06:51:50.68Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/aa/7f/2eaa388267a78401f6e182662b08a588ef4f3de6f0eab1ec09736a7aaa2b/msgpack-1.1.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:36043272c6aede309d29d56851f8841ba907a1a3d04435e43e8a19928e243c1d", size = 79279, upload-time = "2025-06-13T06:51:51.72Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f8/46/31eb60f4452c96161e4dfd26dbca562b4ec68c72e4ad07d9566d7ea35e8a/msgpack-1.1.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a32747b1b39c3ac27d0670122b57e6e57f28eefb725e0b625618d1b59bf9d1e0", size = 423859, upload-time = "2025-06-13T06:51:52.749Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/45/16/a20fa8c32825cc7ae8457fab45670c7a8996d7746ce80ce41cc51e3b2bd7/msgpack-1.1.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8a8b10fdb84a43e50d38057b06901ec9da52baac6983d3f709d8507f3889d43f", size = 429975, upload-time = "2025-06-13T06:51:53.97Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/86/ea/6c958e07692367feeb1a1594d35e22b62f7f476f3c568b002a5ea09d443d/msgpack-1.1.1-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ba0c325c3f485dc54ec298d8b024e134acf07c10d494ffa24373bea729acf704", size = 413528, upload-time = "2025-06-13T06:51:55.507Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/75/05/ac84063c5dae79722bda9f68b878dc31fc3059adb8633c79f1e82c2cd946/msgpack-1.1.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:88daaf7d146e48ec71212ce21109b66e06a98e5e44dca47d853cbfe171d6c8d2", size = 413338, upload-time = "2025-06-13T06:51:57.023Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/69/e8/fe86b082c781d3e1c09ca0f4dacd457ede60a13119b6ce939efe2ea77b76/msgpack-1.1.1-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:d8b55ea20dc59b181d3f47103f113e6f28a5e1c89fd5b67b9140edb442ab67f2", size = 422658, upload-time = "2025-06-13T06:51:58.419Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3b/2b/bafc9924df52d8f3bb7c00d24e57be477f4d0f967c0a31ef5e2225e035c7/msgpack-1.1.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:4a28e8072ae9779f20427af07f53bbb8b4aa81151054e882aee333b158da8752", size = 427124, upload-time = "2025-06-13T06:51:59.969Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a2/3b/1f717e17e53e0ed0b68fa59e9188f3f610c79d7151f0e52ff3cd8eb6b2dc/msgpack-1.1.1-cp311-cp311-win32.whl", hash = "sha256:7da8831f9a0fdb526621ba09a281fadc58ea12701bc709e7b8cbc362feabc295", size = 65016, upload-time = "2025-06-13T06:52:01.294Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/48/45/9d1780768d3b249accecc5a38c725eb1e203d44a191f7b7ff1941f7df60c/msgpack-1.1.1-cp311-cp311-win_amd64.whl", hash = "sha256:5fd1b58e1431008a57247d6e7cc4faa41c3607e8e7d4aaf81f7c29ea013cb458", size = 72267, upload-time = "2025-06-13T06:52:02.568Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e3/26/389b9c593eda2b8551b2e7126ad3a06af6f9b44274eb3a4f054d48ff7e47/msgpack-1.1.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:ae497b11f4c21558d95de9f64fff7053544f4d1a17731c866143ed6bb4591238", size = 82359, upload-time = "2025-06-13T06:52:03.909Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ab/65/7d1de38c8a22cf8b1551469159d4b6cf49be2126adc2482de50976084d78/msgpack-1.1.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:33be9ab121df9b6b461ff91baac6f2731f83d9b27ed948c5b9d1978ae28bf157", size = 79172, upload-time = "2025-06-13T06:52:05.246Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0f/bd/cacf208b64d9577a62c74b677e1ada005caa9b69a05a599889d6fc2ab20a/msgpack-1.1.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6f64ae8fe7ffba251fecb8408540c34ee9df1c26674c50c4544d72dbf792e5ce", size = 425013, upload-time = "2025-06-13T06:52:06.341Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4d/ec/fd869e2567cc9c01278a736cfd1697941ba0d4b81a43e0aa2e8d71dab208/msgpack-1.1.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a494554874691720ba5891c9b0b39474ba43ffb1aaf32a5dac874effb1619e1a", size = 426905, upload-time = "2025-06-13T06:52:07.501Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/55/2a/35860f33229075bce803a5593d046d8b489d7ba2fc85701e714fc1aaf898/msgpack-1.1.1-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:cb643284ab0ed26f6957d969fe0dd8bb17beb567beb8998140b5e38a90974f6c", size = 407336, upload-time = "2025-06-13T06:52:09.047Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8c/16/69ed8f3ada150bf92745fb4921bd621fd2cdf5a42e25eb50bcc57a5328f0/msgpack-1.1.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:d275a9e3c81b1093c060c3837e580c37f47c51eca031f7b5fb76f7b8470f5f9b", size = 409485, upload-time = "2025-06-13T06:52:10.382Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c6/b6/0c398039e4c6d0b2e37c61d7e0e9d13439f91f780686deb8ee64ecf1ae71/msgpack-1.1.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:4fd6b577e4541676e0cc9ddc1709d25014d3ad9a66caa19962c4f5de30fc09ef", size = 412182, upload-time = "2025-06-13T06:52:11.644Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b8/d0/0cf4a6ecb9bc960d624c93effaeaae75cbf00b3bc4a54f35c8507273cda1/msgpack-1.1.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:bb29aaa613c0a1c40d1af111abf025f1732cab333f96f285d6a93b934738a68a", size = 419883, upload-time = "2025-06-13T06:52:12.806Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/62/83/9697c211720fa71a2dfb632cad6196a8af3abea56eece220fde4674dc44b/msgpack-1.1.1-cp312-cp312-win32.whl", hash = "sha256:870b9a626280c86cff9c576ec0d9cbcc54a1e5ebda9cd26dab12baf41fee218c", size = 65406, upload-time = "2025-06-13T06:52:14.271Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c0/23/0abb886e80eab08f5e8c485d6f13924028602829f63b8f5fa25a06636628/msgpack-1.1.1-cp312-cp312-win_amd64.whl", hash = "sha256:5692095123007180dca3e788bb4c399cc26626da51629a31d40207cb262e67f4", size = 72558, upload-time = "2025-06-13T06:52:15.252Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a1/38/561f01cf3577430b59b340b51329803d3a5bf6a45864a55f4ef308ac11e3/msgpack-1.1.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:3765afa6bd4832fc11c3749be4ba4b69a0e8d7b728f78e68120a157a4c5d41f0", size = 81677, upload-time = "2025-06-13T06:52:16.64Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/09/48/54a89579ea36b6ae0ee001cba8c61f776451fad3c9306cd80f5b5c55be87/msgpack-1.1.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:8ddb2bcfd1a8b9e431c8d6f4f7db0773084e107730ecf3472f1dfe9ad583f3d9", size = 78603, upload-time = "2025-06-13T06:52:17.843Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a0/60/daba2699b308e95ae792cdc2ef092a38eb5ee422f9d2fbd4101526d8a210/msgpack-1.1.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:196a736f0526a03653d829d7d4c5500a97eea3648aebfd4b6743875f28aa2af8", size = 420504, upload-time = "2025-06-13T06:52:18.982Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/20/22/2ebae7ae43cd8f2debc35c631172ddf14e2a87ffcc04cf43ff9df9fff0d3/msgpack-1.1.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9d592d06e3cc2f537ceeeb23d38799c6ad83255289bb84c2e5792e5a8dea268a", size = 423749, upload-time = "2025-06-13T06:52:20.211Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/40/1b/54c08dd5452427e1179a40b4b607e37e2664bca1c790c60c442c8e972e47/msgpack-1.1.1-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:4df2311b0ce24f06ba253fda361f938dfecd7b961576f9be3f3fbd60e87130ac", size = 404458, upload-time = "2025-06-13T06:52:21.429Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2e/60/6bb17e9ffb080616a51f09928fdd5cac1353c9becc6c4a8abd4e57269a16/msgpack-1.1.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:e4141c5a32b5e37905b5940aacbc59739f036930367d7acce7a64e4dec1f5e0b", size = 405976, upload-time = "2025-06-13T06:52:22.995Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ee/97/88983e266572e8707c1f4b99c8fd04f9eb97b43f2db40e3172d87d8642db/msgpack-1.1.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:b1ce7f41670c5a69e1389420436f41385b1aa2504c3b0c30620764b15dded2e7", size = 408607, upload-time = "2025-06-13T06:52:24.152Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/bc/66/36c78af2efaffcc15a5a61ae0df53a1d025f2680122e2a9eb8442fed3ae4/msgpack-1.1.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:4147151acabb9caed4e474c3344181e91ff7a388b888f1e19ea04f7e73dc7ad5", size = 424172, upload-time = "2025-06-13T06:52:25.704Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8c/87/a75eb622b555708fe0427fab96056d39d4c9892b0c784b3a721088c7ee37/msgpack-1.1.1-cp313-cp313-win32.whl", hash = "sha256:500e85823a27d6d9bba1d057c871b4210c1dd6fb01fbb764e37e4e8847376323", size = 65347, upload-time = "2025-06-13T06:52:26.846Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ca/91/7dc28d5e2a11a5ad804cf2b7f7a5fcb1eb5a4966d66a5d2b41aee6376543/msgpack-1.1.1-cp313-cp313-win_amd64.whl", hash = "sha256:6d489fba546295983abd142812bda76b57e33d0b9f5d5b71c09a583285506f69", size = 72341, upload-time = "2025-06-13T06:52:27.835Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/bd/74/b0fcaec0cea3f104c61c646f49571864f12321de7b8705e98a32d29ba2ad/msgpack-1.1.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bba1be28247e68994355e028dcd668316db30c1f758d3241a7b903ac78dcd285", size = 409181, upload-time = "2025-06-13T06:52:28.835Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7e/a4/257806f574f8b4bfb76d428b2406cf4585d9f9b582887a0f466278bf0e2a/msgpack-1.1.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b8f93dcddb243159c9e4109c9750ba5b335ab8d48d9522c5308cd05d7e3ce600", size = 413772, upload-time = "2025-06-13T06:52:29.997Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/96/17/46438f4848e86e2f481d46bd3f8b0b0405243b4125bac28ce86dc01e3aeb/msgpack-1.1.1-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:2fbbc0b906a24038c9958a1ba7ae0918ad35b06cb449d398b76a7d08470b0ed9", size = 402772, upload-time = "2025-06-13T06:52:31.195Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1d/72/0ba95da893ddffb09975b4e81fd7b7e612aace0a42ce0d9bdd1a7d802cfe/msgpack-1.1.1-cp38-cp38-musllinux_1_2_aarch64.whl", hash = "sha256:61e35a55a546a1690d9d09effaa436c25ae6130573b6ee9829c37ef0f18d5e78", size = 404650, upload-time = "2025-06-13T06:52:32.638Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/85/d2/c849832b0c0bfb241efc830ccbe7fb880274bbdbc4780798b835f2cd7b3b/msgpack-1.1.1-cp38-cp38-musllinux_1_2_i686.whl", hash = "sha256:1abfc6e949b352dadf4bce0eb78023212ec5ac42f6abfd469ce91d783c149c2a", size = 413595, upload-time = "2025-06-13T06:52:33.882Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/03/79/ea7cda493ec78afb9bd4c88e3c8bf5bffabca78d1917d8b24cddd0b9f5ee/msgpack-1.1.1-cp38-cp38-musllinux_1_2_x86_64.whl", hash = "sha256:996f2609ddf0142daba4cefd767d6db26958aac8439ee41db9cc0db9f4c4c3a6", size = 412830, upload-time = "2025-06-13T06:52:35.431Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e3/80/644311ca3064cfc9a9ecf64074e905e5359da730faefc88c6cfbbaf110ee/msgpack-1.1.1-cp38-cp38-win32.whl", hash = "sha256:4d3237b224b930d58e9d83c81c0dba7aacc20fcc2f89c1e5423aa0529a4cd142", size = 65439, upload-time = "2025-06-13T06:52:36.976Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ce/ec/27d4740fdeea71a7d559b405614b5d9b866028768a949e8dd58abed8474f/msgpack-1.1.1-cp38-cp38-win_amd64.whl", hash = "sha256:da8f41e602574ece93dbbda1fab24650d6bf2a24089f9e9dbb4f5730ec1e58ad", size = 72234, upload-time = "2025-06-13T06:52:37.966Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1f/bd/0792be119d7fe7dc2148689ef65c90507d82d20a204aab3b98c74a1f8684/msgpack-1.1.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:f5be6b6bc52fad84d010cb45433720327ce886009d862f46b26d4d154001994b", size = 81882, upload-time = "2025-06-13T06:52:39.316Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/75/77/ce06c8e26a816ae8730a8e030d263c5289adcaff9f0476f9b270bdd7c5c2/msgpack-1.1.1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:3a89cd8c087ea67e64844287ea52888239cbd2940884eafd2dcd25754fb72232", size = 78414, upload-time = "2025-06-13T06:52:40.341Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/73/27/190576c497677fb4a0d05d896b24aea6cdccd910f206aaa7b511901befed/msgpack-1.1.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1d75f3807a9900a7d575d8d6674a3a47e9f227e8716256f35bc6f03fc597ffbf", size = 400927, upload-time = "2025-06-13T06:52:41.399Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ed/af/6a0aa5a06762e70726ec3c10fb966600d84a7220b52635cb0ab2dc64d32f/msgpack-1.1.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d182dac0221eb8faef2e6f44701812b467c02674a322c739355c39e94730cdbf", size = 405903, upload-time = "2025-06-13T06:52:42.699Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1e/80/3f3da358cecbbe8eb12360814bd1277d59d2608485934742a074d99894a9/msgpack-1.1.1-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1b13fe0fb4aac1aa5320cd693b297fe6fdef0e7bea5518cbc2dd5299f873ae90", size = 393192, upload-time = "2025-06-13T06:52:43.986Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/98/c6/3a0ec7fdebbb4f3f8f254696cd91d491c29c501dbebd86286c17e8f68cd7/msgpack-1.1.1-cp39-cp39-musllinux_1_2_aarch64.whl", hash = "sha256:435807eeb1bc791ceb3247d13c79868deb22184e1fc4224808750f0d7d1affc1", size = 393851, upload-time = "2025-06-13T06:52:45.177Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/39/37/df50d5f8e68514b60fbe70f6e8337ea2b32ae2be030871bcd9d1cf7d4b62/msgpack-1.1.1-cp39-cp39-musllinux_1_2_i686.whl", hash = "sha256:4835d17af722609a45e16037bb1d4d78b7bdf19d6c0128116d178956618c4e88", size = 400292, upload-time = "2025-06-13T06:52:46.381Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fc/ec/1e067292e02d2ceb4c8cb5ba222c4f7bb28730eef5676740609dc2627e0f/msgpack-1.1.1-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:a8ef6e342c137888ebbfb233e02b8fbd689bb5b5fcc59b34711ac47ebd504478", size = 401873, upload-time = "2025-06-13T06:52:47.957Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d3/31/e8c9c6b5b58d64c9efa99c8d181fcc25f38ead357b0360379fbc8a4234ad/msgpack-1.1.1-cp39-cp39-win32.whl", hash = "sha256:61abccf9de335d9efd149e2fff97ed5974f2481b3353772e8e2dd3402ba2bd57", size = 65028, upload-time = "2025-06-13T06:52:49.166Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/20/d6/cd62cded572e5e25892747a5d27850170bcd03c855e9c69c538e024de6f9/msgpack-1.1.1-cp39-cp39-win_amd64.whl", hash = "sha256:40eae974c873b2992fd36424a5d9407f93e97656d999f43fca9d29f820899084", size = 71700, upload-time = "2025-06-13T06:52:50.244Z" }, +] + +[[package]] +name = "networkx" +version = "3.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fd/a1/47b974da1a73f063c158a1f4cc33ed0abf7c04f98a19050e80c533c31f0c/networkx-3.1.tar.gz", hash = "sha256:de346335408f84de0eada6ff9fafafff9bcda11f0a0dfaa931133debb146ab61", size = 2021691, upload-time = "2023-04-04T20:07:56.693Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a8/05/9d4f9b78ead6b2661d6e8ea772e111fc4a9fbd866ad0c81906c11206b55e/networkx-3.1-py3-none-any.whl", hash = "sha256:4f33f68cb2afcf86f28a45f43efc27a9386b535d567d2127f8f61d51dec58d36", size = 2072251, upload-time = "2023-04-04T20:07:53.63Z" }, +] + +[[package]] +name = "networkx" +version = "3.2.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform == 'darwin')", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c4/80/a84676339aaae2f1cfdf9f418701dd634aef9cc76f708ef55c36ff39c3ca/networkx-3.2.1.tar.gz", hash = "sha256:9f1bb5cf3409bf324e0a722c20bdb4c20ee39bf1c30ce8ae499c8502b0b5e0c6", size = 2073928, upload-time = "2023-10-28T08:41:39.364Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d5/f0/8fbc882ca80cf077f1b246c0e3c3465f7f415439bdea6b899f6b19f61f70/networkx-3.2.1-py3-none-any.whl", hash = "sha256:f18c69adc97877c42332c170849c96cefa91881c99a7cb3e95b7c659ebdc1ec2", size = 1647772, upload-time = "2023-10-28T08:41:36.945Z" }, +] + +[[package]] +name = "networkx" +version = "3.4.2" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'darwin')", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fd/1d/06475e1cd5264c0b870ea2cc6fdb3e37177c1e565c43f56ff17a10e3937f/networkx-3.4.2.tar.gz", hash = "sha256:307c3669428c5362aab27c8a1260aa8f47c4e91d3891f48be0141738d8d053e1", size = 2151368, upload-time = "2024-10-21T12:39:38.695Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b9/54/dd730b32ea14ea797530a4479b2ed46a6fb250f682a9cfb997e968bf0261/networkx-3.4.2-py3-none-any.whl", hash = "sha256:df5d4365b724cf81b8c6a7312509d0c22386097011ad1abe274afd5e9d3bbc5f", size = 1723263, upload-time = "2024-10-21T12:39:36.247Z" }, +] + +[[package]] +name = "networkx" +version = "3.5" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.13.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'darwin')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'darwin')", + "(python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform == 'darwin')", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6c/4f/ccdb8ad3a38e583f214547fd2f7ff1fc160c43a75af88e6aec213404b96a/networkx-3.5.tar.gz", hash = "sha256:d4c6f9cf81f52d69230866796b82afbccdec3db7ae4fbd1b65ea750feed50037", size = 2471065, upload-time = "2025-05-29T11:35:07.804Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/eb/8d/776adee7bbf76365fdd7f2552710282c79a4ead5d2a46408c9043a2b70ba/networkx-3.5-py3-none-any.whl", hash = "sha256:0030d386a9a06dee3565298b4a734b68589749a544acbb6c412dc9e2489ec6ec", size = 2034406, upload-time = "2025-05-29T11:35:04.961Z" }, +] + +[[package]] +name = "numba" +version = "0.58.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] +dependencies = [ + { name = "importlib-metadata", version = "8.5.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "llvmlite", version = "0.41.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "numpy", version = "1.24.4", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/13/2b/0f750d451fd961fd91d6bc86c512781fa46f9ef64813007e501482522ff9/numba-0.58.1.tar.gz", hash = "sha256:487ded0633efccd9ca3a46364b40006dbdaca0f95e99b8b83e778d1195ebcbaa", size = 2623830, upload-time = "2023-10-18T13:47:46.54Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8d/e6/7e2c73fcdf3c98dcc0d3035ecafab214ed9f1978ca18a149bad6416bd308/numba-0.58.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:07f2fa7e7144aa6f275f27260e73ce0d808d3c62b30cff8906ad1dec12d87bbe", size = 2588424, upload-time = "2023-10-18T13:46:37.554Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/39/44/e5540c376f6b99d906f46bdf98ce4d91511a0ee89ce60a62e2498dacd6fe/numba-0.58.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:7bf1ddd4f7b9c2306de0384bf3854cac3edd7b4d8dffae2ec1b925e4c436233f", size = 2591283, upload-time = "2023-10-18T13:46:40.19Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/93/fa/40cf7aff6f5836412d8f5f751a92be3efd249ed322a01a5c9c1fe4c4a7a9/numba-0.58.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:bc2d904d0319d7a5857bd65062340bed627f5bfe9ae4a495aef342f072880d50", size = 3349802, upload-time = "2023-10-18T13:46:42.688Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ed/13/b66627125b35f2987bd9872cf028b5e1e1ffcbc8d1e182ac4e84eed3998f/numba-0.58.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4e79b6cc0d2bf064a955934a2e02bf676bc7995ab2db929dbbc62e4c16551be6", size = 3641467, upload-time = "2023-10-18T13:46:45.099Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3b/77/fdf609595d4f8f7f7e6e0fcaf11f7223d636cae4cb5c0805f46778c08a57/numba-0.58.1-cp310-cp310-win_amd64.whl", hash = "sha256:81fe5b51532478149b5081311b0fd4206959174e660c372b94ed5364cfb37c82", size = 2606344, upload-time = "2023-10-18T13:46:47.204Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/9a/ce/a1e69316135646e83bc77b9eb662e3dc9414a6fea9eef2db4fd4c9fcf5e3/numba-0.58.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:bcecd3fb9df36554b342140a4d77d938a549be635d64caf8bd9ef6c47a47f8aa", size = 2588116, upload-time = "2023-10-18T13:46:49.385Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/88/60/aa67255e6e166ef6541d007f22470fc459b8c7b66fb6790fd735d0bcd951/numba-0.58.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:a1eaa744f518bbd60e1f7ccddfb8002b3d06bd865b94a5d7eac25028efe0e0ff", size = 2590990, upload-time = "2023-10-18T13:46:51.646Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ca/f8/18f24441783f345a805a08bdcfa7cfb9c987a6154da412c734ed04489f07/numba-0.58.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:bf68df9c307fb0aa81cacd33faccd6e419496fdc621e83f1efce35cdc5e79cac", size = 3355149, upload-time = "2023-10-18T13:46:54.537Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/65/85/72649823bcbfd01820e1c60fa54eb77adb2dba8be921fe6fabdbeb011a88/numba-0.58.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:55a01e1881120e86d54efdff1be08381886fe9f04fc3006af309c602a72bc44d", size = 3647131, upload-time = "2023-10-18T13:46:57.491Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/cd/59/5dd8a3059997ec1daf6f9f7c9b1aef7f0a9e23e1334a5774eae65cae6fd0/numba-0.58.1-cp311-cp311-win_amd64.whl", hash = "sha256:811305d5dc40ae43c3ace5b192c670c358a89a4d2ae4f86d1665003798ea7a1a", size = 2606316, upload-time = "2023-10-18T13:47:03.176Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/76/ff/9e6726d73effefc14557c869c980898f123749c16fae4c527285169f1d9c/numba-0.58.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:ea5bfcf7d641d351c6a80e8e1826eb4a145d619870016eeaf20bbd71ef5caa22", size = 2584254, upload-time = "2023-10-18T13:47:07.011Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fa/01/b8b3d3c4623f3f267d7f52f30f2f994cc00793e67a6242d0b7323099060e/numba-0.58.1-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:e63d6aacaae1ba4ef3695f1c2122b30fa3d8ba039c8f517784668075856d79e2", size = 2589312, upload-time = "2023-10-18T13:47:09.887Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c6/97/dad816d9313041aa967e13f3b21fdff040433c7efe9eff3112bcf3b40bc1/numba-0.58.1-cp38-cp38-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6fe7a9d8e3bd996fbe5eac0683227ccef26cba98dae6e5cee2c1894d4b9f16c1", size = 3359774, upload-time = "2023-10-18T13:47:13.8Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6e/03/c658004b4b4d52dca7a74582e9fc3e20f94c17cb4b797eb3601876759540/numba-0.58.1-cp38-cp38-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:898af055b03f09d33a587e9425500e5be84fc90cd2f80b3fb71c6a4a17a7e354", size = 3651491, upload-time = "2023-10-18T13:47:18.489Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ce/76/9266f7bd78f134f6ae09886d87675a8ac34e8dd0a30f6007794da41d8a80/numba-0.58.1-cp38-cp38-win_amd64.whl", hash = "sha256:d3e2fe81fe9a59fcd99cc572002101119059d64d31eb6324995ee8b0f144a306", size = 2626200, upload-time = "2023-10-18T13:47:23.899Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b5/de/e2ef933a99c502d2ec5dda1a43a74ab98b1b606c0ff17422d42c62a6f00f/numba-0.58.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:5c765aef472a9406a97ea9782116335ad4f9ef5c9f93fc05fd44aab0db486954", size = 2587997, upload-time = "2023-10-18T13:47:27.885Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/eb/61/db566ad7cd3f7e9191ec744f25cf6455e8d91732ad70be8fcdd203a48dab/numba-0.58.1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:9e9356e943617f5e35a74bf56ff6e7cc83e6b1865d5e13cee535d79bf2cae954", size = 2590940, upload-time = "2023-10-18T13:47:31.597Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3d/db/75bd7bb82e0f05ad066a0f4b1f8913de771dcdec37b063caa4c9282e6b1b/numba-0.58.1-cp39-cp39-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:240e7a1ae80eb6b14061dc91263b99dc8d6af9ea45d310751b780888097c1aaa", size = 3345620, upload-time = "2023-10-18T13:47:35.316Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8b/4b/58bdecd06266c5395157727d7191d83c82d190665bdf3d65fdd67a6dbaa5/numba-0.58.1-cp39-cp39-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:45698b995914003f890ad839cfc909eeb9c74921849c712a05405d1a79c50f68", size = 3636967, upload-time = "2023-10-18T13:47:38.751Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f8/14/546925a397d7638858a890a01d82cc6de147fca294eccf93ee11b524ca3c/numba-0.58.1-cp39-cp39-win_amd64.whl", hash = "sha256:bd3dda77955be03ff366eebbfdb39919ce7c2620d86c906203bed92124989032", size = 2626240, upload-time = "2023-10-18T13:47:42.379Z" }, +] + +[[package]] +name = "numba" +version = "0.60.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform == 'darwin')", +] +dependencies = [ + { name = "llvmlite", version = "0.43.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, + { name = "numpy", version = "2.0.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3c/93/2849300a9184775ba274aba6f82f303343669b0592b7bb0849ea713dabb0/numba-0.60.0.tar.gz", hash = "sha256:5df6158e5584eece5fc83294b949fd30b9f1125df7708862205217e068aabf16", size = 2702171, upload-time = "2024-06-13T18:11:19.869Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f7/cf/baa13a7e3556d73d9e38021e6d6aa4aeb30d8b94545aa8b70d0f24a1ccc4/numba-0.60.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:5d761de835cd38fb400d2c26bb103a2726f548dc30368853121d66201672e651", size = 2647627, upload-time = "2024-06-13T18:10:29.857Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ac/ba/4b57fa498564457c3cc9fc9e570a6b08e6086c74220f24baaf04e54b995f/numba-0.60.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:159e618ef213fba758837f9837fb402bbe65326e60ba0633dbe6c7f274d42c1b", size = 2650322, upload-time = "2024-06-13T18:10:32.849Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/28/98/7ea97ee75870a54f938a8c70f7e0be4495ba5349c5f9db09d467c4a5d5b7/numba-0.60.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:1527dc578b95c7c4ff248792ec33d097ba6bef9eda466c948b68dfc995c25781", size = 3407390, upload-time = "2024-06-13T18:10:34.741Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/79/58/cb4ac5b8f7ec64200460aef1fed88258fb872ceef504ab1f989d2ff0f684/numba-0.60.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fe0b28abb8d70f8160798f4de9d486143200f34458d34c4a214114e445d7124e", size = 3699694, upload-time = "2024-06-13T18:10:37.295Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1c/b0/c61a93ca947d12233ff45de506ddbf52af3f752066a0b8be4d27426e16da/numba-0.60.0-cp310-cp310-win_amd64.whl", hash = "sha256:19407ced081d7e2e4b8d8c36aa57b7452e0283871c296e12d798852bc7d7f198", size = 2687030, upload-time = "2024-06-13T18:10:39.47Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/98/ad/df18d492a8f00d29a30db307904b9b296e37507034eedb523876f3a2e13e/numba-0.60.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:a17b70fc9e380ee29c42717e8cc0bfaa5556c416d94f9aa96ba13acb41bdece8", size = 2647254, upload-time = "2024-06-13T18:10:41.69Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/9a/51/a4dc2c01ce7a850b8e56ff6d5381d047a5daea83d12bad08aa071d34b2ee/numba-0.60.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:3fb02b344a2a80efa6f677aa5c40cd5dd452e1b35f8d1c2af0dfd9ada9978e4b", size = 2649970, upload-time = "2024-06-13T18:10:44.682Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f9/4c/8889ac94c0b33dca80bed11564b8c6d9ea14d7f094e674c58e5c5b05859b/numba-0.60.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:5f4fde652ea604ea3c86508a3fb31556a6157b2c76c8b51b1d45eb40c8598703", size = 3412492, upload-time = "2024-06-13T18:10:47.1Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/57/03/2b4245b05b71c0cee667e6a0b51606dfa7f4157c9093d71c6b208385a611/numba-0.60.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4142d7ac0210cc86432b818338a2bc368dc773a2f5cf1e32ff7c5b378bd63ee8", size = 3705018, upload-time = "2024-06-13T18:10:49.539Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/79/89/2d924ca60dbf949f18a6fec223a2445f5f428d9a5f97a6b29c2122319015/numba-0.60.0-cp311-cp311-win_amd64.whl", hash = "sha256:cac02c041e9b5bc8cf8f2034ff6f0dbafccd1ae9590dc146b3a02a45e53af4e2", size = 2686920, upload-time = "2024-06-13T18:10:51.937Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/eb/5c/b5ec752c475e78a6c3676b67c514220dbde2725896bbb0b6ec6ea54b2738/numba-0.60.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:d7da4098db31182fc5ffe4bc42c6f24cd7d1cb8a14b59fd755bfee32e34b8404", size = 2647866, upload-time = "2024-06-13T18:10:54.453Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/65/42/39559664b2e7c15689a638c2a38b3b74c6e69a04e2b3019b9f7742479188/numba-0.60.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:38d6ea4c1f56417076ecf8fc327c831ae793282e0ff51080c5094cb726507b1c", size = 2650208, upload-time = "2024-06-13T18:10:56.779Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/67/88/c4459ccc05674ef02119abf2888ccd3e2fed12a323f52255f4982fc95876/numba-0.60.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:62908d29fb6a3229c242e981ca27e32a6e606cc253fc9e8faeb0e48760de241e", size = 3466946, upload-time = "2024-06-13T18:10:58.961Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8b/41/ac11cf33524def12aa5bd698226ae196a1185831c05ed29dc0c56eaa308b/numba-0.60.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:0ebaa91538e996f708f1ab30ef4d3ddc344b64b5227b67a57aa74f401bb68b9d", size = 3761463, upload-time = "2024-06-13T18:11:01.657Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ca/bd/0fe29fcd1b6a8de479a4ed25c6e56470e467e3611c079d55869ceef2b6d1/numba-0.60.0-cp312-cp312-win_amd64.whl", hash = "sha256:f75262e8fe7fa96db1dca93d53a194a38c46da28b112b8a4aca168f0df860347", size = 2707588, upload-time = "2024-06-13T18:11:04.261Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/68/1a/87c53f836cdf557083248c3f47212271f220280ff766538795e77c8c6bbf/numba-0.60.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:01ef4cd7d83abe087d644eaa3d95831b777aa21d441a23703d649e06b8e06b74", size = 2647186, upload-time = "2024-06-13T18:11:06.753Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/28/14/a5baa1f2edea7b49afa4dc1bb1b126645198cf1075186853b5b497be826e/numba-0.60.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:819a3dfd4630d95fd574036f99e47212a1af41cbcb019bf8afac63ff56834449", size = 2650038, upload-time = "2024-06-13T18:11:10.869Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3b/bd/f1985719ff34e37e07bb18f9d3acd17e5a21da255f550c8eae031e2ddf5f/numba-0.60.0-cp39-cp39-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:0b983bd6ad82fe868493012487f34eae8bf7dd94654951404114f23c3466d34b", size = 3403010, upload-time = "2024-06-13T18:11:13.057Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/54/9b/cd73d3f6617ddc8398a63ef97d8dc9139a9879b9ca8a7ca4b8789056ea46/numba-0.60.0-cp39-cp39-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c151748cd269ddeab66334bd754817ffc0cabd9433acb0f551697e5151917d25", size = 3695086, upload-time = "2024-06-13T18:11:15.497Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/01/01/8b7b670c77c5ea0e47e283d82332969bf672ab6410d0b2610cac5b7a3ded/numba-0.60.0-cp39-cp39-win_amd64.whl", hash = "sha256:3031547a015710140e8c87226b4cfe927cac199835e5bf7d4fe5cb64e814e3ab", size = 2686978, upload-time = "2024-06-13T18:11:17.765Z" }, +] + +[[package]] +name = "numba" +version = "0.61.2" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.13.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'darwin')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'darwin')", + "(python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform == 'darwin')", + "(python_full_version == '3.10.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'darwin')", +] +dependencies = [ + { name = "llvmlite", version = "0.44.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.10'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.10'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1c/a0/e21f57604304aa03ebb8e098429222722ad99176a4f979d34af1d1ee80da/numba-0.61.2.tar.gz", hash = "sha256:8750ee147940a6637b80ecf7f95062185ad8726c8c28a2295b8ec1160a196f7d", size = 2820615, upload-time = "2025-04-09T02:58:07.659Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/eb/ca/f470be59552ccbf9531d2d383b67ae0b9b524d435fb4a0d229fef135116e/numba-0.61.2-cp310-cp310-macosx_10_14_x86_64.whl", hash = "sha256:cf9f9fc00d6eca0c23fc840817ce9f439b9f03c8f03d6246c0e7f0cb15b7162a", size = 2775663, upload-time = "2025-04-09T02:57:34.143Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f5/13/3bdf52609c80d460a3b4acfb9fdb3817e392875c0d6270cf3fd9546f138b/numba-0.61.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:ea0247617edcb5dd61f6106a56255baab031acc4257bddaeddb3a1003b4ca3fd", size = 2778344, upload-time = "2025-04-09T02:57:36.609Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e2/7d/bfb2805bcfbd479f04f835241ecf28519f6e3609912e3a985aed45e21370/numba-0.61.2-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ae8c7a522c26215d5f62ebec436e3d341f7f590079245a2f1008dfd498cc1642", size = 3824054, upload-time = "2025-04-09T02:57:38.162Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e3/27/797b2004745c92955470c73c82f0e300cf033c791f45bdecb4b33b12bdea/numba-0.61.2-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:bd1e74609855aa43661edffca37346e4e8462f6903889917e9f41db40907daa2", size = 3518531, upload-time = "2025-04-09T02:57:39.709Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b1/c6/c2fb11e50482cb310afae87a997707f6c7d8a48967b9696271347441f650/numba-0.61.2-cp310-cp310-win_amd64.whl", hash = "sha256:ae45830b129c6137294093b269ef0a22998ccc27bf7cf096ab8dcf7bca8946f9", size = 2831612, upload-time = "2025-04-09T02:57:41.559Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3f/97/c99d1056aed767503c228f7099dc11c402906b42a4757fec2819329abb98/numba-0.61.2-cp311-cp311-macosx_10_14_x86_64.whl", hash = "sha256:efd3db391df53aaa5cfbee189b6c910a5b471488749fd6606c3f33fc984c2ae2", size = 2775825, upload-time = "2025-04-09T02:57:43.442Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/95/9e/63c549f37136e892f006260c3e2613d09d5120672378191f2dc387ba65a2/numba-0.61.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:49c980e4171948ffebf6b9a2520ea81feed113c1f4890747ba7f59e74be84b1b", size = 2778695, upload-time = "2025-04-09T02:57:44.968Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/97/c8/8740616c8436c86c1b9a62e72cb891177d2c34c2d24ddcde4c390371bf4c/numba-0.61.2-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3945615cd73c2c7eba2a85ccc9c1730c21cd3958bfcf5a44302abae0fb07bb60", size = 3829227, upload-time = "2025-04-09T02:57:46.63Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fc/06/66e99ae06507c31d15ff3ecd1f108f2f59e18b6e08662cd5f8a5853fbd18/numba-0.61.2-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:bbfdf4eca202cebade0b7d43896978e146f39398909a42941c9303f82f403a18", size = 3523422, upload-time = "2025-04-09T02:57:48.222Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0f/a4/2b309a6a9f6d4d8cfba583401c7c2f9ff887adb5d54d8e2e130274c0973f/numba-0.61.2-cp311-cp311-win_amd64.whl", hash = "sha256:76bcec9f46259cedf888041b9886e257ae101c6268261b19fda8cfbc52bec9d1", size = 2831505, upload-time = "2025-04-09T02:57:50.108Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b4/a0/c6b7b9c615cfa3b98c4c63f4316e3f6b3bbe2387740277006551784218cd/numba-0.61.2-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:34fba9406078bac7ab052efbf0d13939426c753ad72946baaa5bf9ae0ebb8dd2", size = 2776626, upload-time = "2025-04-09T02:57:51.857Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/92/4a/fe4e3c2ecad72d88f5f8cd04e7f7cff49e718398a2fac02d2947480a00ca/numba-0.61.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:4ddce10009bc097b080fc96876d14c051cc0c7679e99de3e0af59014dab7dfe8", size = 2779287, upload-time = "2025-04-09T02:57:53.658Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/9a/2d/e518df036feab381c23a624dac47f8445ac55686ec7f11083655eb707da3/numba-0.61.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b1bb509d01f23d70325d3a5a0e237cbc9544dd50e50588bc581ba860c213546", size = 3885928, upload-time = "2025-04-09T02:57:55.206Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/10/0f/23cced68ead67b75d77cfcca3df4991d1855c897ee0ff3fe25a56ed82108/numba-0.61.2-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:48a53a3de8f8793526cbe330f2a39fe9a6638efcbf11bd63f3d2f9757ae345cd", size = 3577115, upload-time = "2025-04-09T02:57:56.818Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/68/1d/ddb3e704c5a8fb90142bf9dc195c27db02a08a99f037395503bfbc1d14b3/numba-0.61.2-cp312-cp312-win_amd64.whl", hash = "sha256:97cf4f12c728cf77c9c1d7c23707e4d8fb4632b46275f8f3397de33e5877af18", size = 2831929, upload-time = "2025-04-09T02:57:58.45Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0b/f3/0fe4c1b1f2569e8a18ad90c159298d862f96c3964392a20d74fc628aee44/numba-0.61.2-cp313-cp313-macosx_10_14_x86_64.whl", hash = "sha256:3a10a8fc9afac40b1eac55717cece1b8b1ac0b946f5065c89e00bde646b5b154", size = 2771785, upload-time = "2025-04-09T02:57:59.96Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e9/71/91b277d712e46bd5059f8a5866862ed1116091a7cb03bd2704ba8ebe015f/numba-0.61.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:7d3bcada3c9afba3bed413fba45845f2fb9cd0d2b27dd58a1be90257e293d140", size = 2773289, upload-time = "2025-04-09T02:58:01.435Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0d/e0/5ea04e7ad2c39288c0f0f9e8d47638ad70f28e275d092733b5817cf243c9/numba-0.61.2-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:bdbca73ad81fa196bd53dc12e3aaf1564ae036e0c125f237c7644fe64a4928ab", size = 3893918, upload-time = "2025-04-09T02:58:02.933Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/17/58/064f4dcb7d7e9412f16ecf80ed753f92297e39f399c905389688cf950b81/numba-0.61.2-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:5f154aaea625fb32cfbe3b80c5456d514d416fcdf79733dd69c0df3a11348e9e", size = 3584056, upload-time = "2025-04-09T02:58:04.538Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/af/a4/6d3a0f2d3989e62a18749e1e9913d5fa4910bbb3e3311a035baea6caf26d/numba-0.61.2-cp313-cp313-win_amd64.whl", hash = "sha256:59321215e2e0ac5fa928a8020ab00b8e57cda8a97384963ac0dfa4d4e6aa54e7", size = 2831846, upload-time = "2025-04-09T02:58:06.125Z" }, +] + +[[package]] +name = "numpy" +version = "1.24.4" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a4/9b/027bec52c633f6556dba6b722d9a0befb40498b9ceddd29cbe67a45a127c/numpy-1.24.4.tar.gz", hash = "sha256:80f5e3a4e498641401868df4208b74581206afbee7cf7b8329daae82676d9463", size = 10911229, upload-time = "2023-06-26T13:39:33.218Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6b/80/6cdfb3e275d95155a34659163b83c09e3a3ff9f1456880bec6cc63d71083/numpy-1.24.4-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:c0bfb52d2169d58c1cdb8cc1f16989101639b34c7d3ce60ed70b19c63eba0b64", size = 19789140, upload-time = "2023-06-26T13:22:33.184Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/64/5f/3f01d753e2175cfade1013eea08db99ba1ee4bdb147ebcf3623b75d12aa7/numpy-1.24.4-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:ed094d4f0c177b1b8e7aa9cba7d6ceed51c0e569a5318ac0ca9a090680a6a1b1", size = 13854297, upload-time = "2023-06-26T13:22:59.541Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5a/b3/2f9c21d799fa07053ffa151faccdceeb69beec5a010576b8991f614021f7/numpy-1.24.4-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:79fc682a374c4a8ed08b331bef9c5f582585d1048fa6d80bc6c35bc384eee9b4", size = 13995611, upload-time = "2023-06-26T13:23:22.167Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/10/be/ae5bf4737cb79ba437879915791f6f26d92583c738d7d960ad94e5c36adf/numpy-1.24.4-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7ffe43c74893dbf38c2b0a1f5428760a1a9c98285553c89e12d70a96a7f3a4d6", size = 17282357, upload-time = "2023-06-26T13:23:51.446Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c0/64/908c1087be6285f40e4b3e79454552a701664a079321cff519d8c7051d06/numpy-1.24.4-cp310-cp310-win32.whl", hash = "sha256:4c21decb6ea94057331e111a5bed9a79d335658c27ce2adb580fb4d54f2ad9bc", size = 12429222, upload-time = "2023-06-26T13:24:13.849Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/22/55/3d5a7c1142e0d9329ad27cece17933b0e2ab4e54ddc5c1861fbfeb3f7693/numpy-1.24.4-cp310-cp310-win_amd64.whl", hash = "sha256:b4bea75e47d9586d31e892a7401f76e909712a0fd510f58f5337bea9572c571e", size = 14841514, upload-time = "2023-06-26T13:24:38.129Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a9/cc/5ed2280a27e5dab12994c884f1f4d8c3bd4d885d02ae9e52a9d213a6a5e2/numpy-1.24.4-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:f136bab9c2cfd8da131132c2cf6cc27331dd6fae65f95f69dcd4ae3c3639c810", size = 19775508, upload-time = "2023-06-26T13:25:08.882Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c0/bc/77635c657a3668cf652806210b8662e1aff84b818a55ba88257abf6637a8/numpy-1.24.4-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:e2926dac25b313635e4d6cf4dc4e51c8c0ebfed60b801c799ffc4c32bf3d1254", size = 13840033, upload-time = "2023-06-26T13:25:33.417Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a7/4c/96cdaa34f54c05e97c1c50f39f98d608f96f0677a6589e64e53104e22904/numpy-1.24.4-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:222e40d0e2548690405b0b3c7b21d1169117391c2e82c378467ef9ab4c8f0da7", size = 13991951, upload-time = "2023-06-26T13:25:55.725Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/22/97/dfb1a31bb46686f09e68ea6ac5c63fdee0d22d7b23b8f3f7ea07712869ef/numpy-1.24.4-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7215847ce88a85ce39baf9e89070cb860c98fdddacbaa6c0da3ffb31b3350bd5", size = 17278923, upload-time = "2023-06-26T13:26:25.658Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/35/e2/76a11e54139654a324d107da1d98f99e7aa2a7ef97cfd7c631fba7dbde71/numpy-1.24.4-cp311-cp311-win32.whl", hash = "sha256:4979217d7de511a8d57f4b4b5b2b965f707768440c17cb70fbf254c4b225238d", size = 12422446, upload-time = "2023-06-26T13:26:49.302Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d8/ec/ebef2f7d7c28503f958f0f8b992e7ce606fb74f9e891199329d5f5f87404/numpy-1.24.4-cp311-cp311-win_amd64.whl", hash = "sha256:b7b1fc9864d7d39e28f41d089bfd6353cb5f27ecd9905348c24187a768c79694", size = 14834466, upload-time = "2023-06-26T13:27:16.029Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/11/10/943cfb579f1a02909ff96464c69893b1d25be3731b5d3652c2e0cf1281ea/numpy-1.24.4-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:1452241c290f3e2a312c137a9999cdbf63f78864d63c79039bda65ee86943f61", size = 19780722, upload-time = "2023-06-26T13:27:49.573Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a7/ae/f53b7b265fdc701e663fbb322a8e9d4b14d9cb7b2385f45ddfabfc4327e4/numpy-1.24.4-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:04640dab83f7c6c85abf9cd729c5b65f1ebd0ccf9de90b270cd61935eef0197f", size = 13843102, upload-time = "2023-06-26T13:28:12.288Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/25/6f/2586a50ad72e8dbb1d8381f837008a0321a3516dfd7cb57fc8cf7e4bb06b/numpy-1.24.4-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a5425b114831d1e77e4b5d812b69d11d962e104095a5b9c3b641a218abcc050e", size = 14039616, upload-time = "2023-06-26T13:28:35.659Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/98/5d/5738903efe0ecb73e51eb44feafba32bdba2081263d40c5043568ff60faf/numpy-1.24.4-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dd80e219fd4c71fc3699fc1dadac5dcf4fd882bfc6f7ec53d30fa197b8ee22dc", size = 17316263, upload-time = "2023-06-26T13:29:09.272Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d1/57/8d328f0b91c733aa9aa7ee540dbc49b58796c862b4fbcb1146c701e888da/numpy-1.24.4-cp38-cp38-win32.whl", hash = "sha256:4602244f345453db537be5314d3983dbf5834a9701b7723ec28923e2889e0bb2", size = 12455660, upload-time = "2023-06-26T13:29:33.434Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/69/65/0d47953afa0ad569d12de5f65d964321c208492064c38fe3b0b9744f8d44/numpy-1.24.4-cp38-cp38-win_amd64.whl", hash = "sha256:692f2e0f55794943c5bfff12b3f56f99af76f902fc47487bdfe97856de51a706", size = 14868112, upload-time = "2023-06-26T13:29:58.385Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/9a/cd/d5b0402b801c8a8b56b04c1e85c6165efab298d2f0ab741c2406516ede3a/numpy-1.24.4-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:2541312fbf09977f3b3ad449c4e5f4bb55d0dbf79226d7724211acc905049400", size = 19816549, upload-time = "2023-06-26T13:30:36.976Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/14/27/638aaa446f39113a3ed38b37a66243e21b38110d021bfcb940c383e120f2/numpy-1.24.4-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:9667575fb6d13c95f1b36aca12c5ee3356bf001b714fc354eb5465ce1609e62f", size = 13879950, upload-time = "2023-06-26T13:31:01.787Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8f/27/91894916e50627476cff1a4e4363ab6179d01077d71b9afed41d9e1f18bf/numpy-1.24.4-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f3a86ed21e4f87050382c7bc96571755193c4c1392490744ac73d660e8f564a9", size = 14030228, upload-time = "2023-06-26T13:31:26.696Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7a/7c/d7b2a0417af6428440c0ad7cb9799073e507b1a465f827d058b826236964/numpy-1.24.4-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d11efb4dbecbdf22508d55e48d9c8384db795e1b7b51ea735289ff96613ff74d", size = 17311170, upload-time = "2023-06-26T13:31:56.615Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/18/9d/e02ace5d7dfccee796c37b995c63322674daf88ae2f4a4724c5dd0afcc91/numpy-1.24.4-cp39-cp39-win32.whl", hash = "sha256:6620c0acd41dbcb368610bb2f4d83145674040025e5536954782467100aa8835", size = 12454918, upload-time = "2023-06-26T13:32:16.8Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/63/38/6cc19d6b8bfa1d1a459daf2b3fe325453153ca7019976274b6f33d8b5663/numpy-1.24.4-cp39-cp39-win_amd64.whl", hash = "sha256:befe2bf740fd8373cf56149a5c23a0f601e82869598d41f8e188a0e9869926f8", size = 14867441, upload-time = "2023-06-26T13:32:40.521Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a4/fd/8dff40e25e937c94257455c237b9b6bf5a30d42dd1cc11555533be099492/numpy-1.24.4-pp38-pypy38_pp73-macosx_10_9_x86_64.whl", hash = "sha256:31f13e25b4e304632a4619d0e0777662c2ffea99fcae2029556b17d8ff958aef", size = 19156590, upload-time = "2023-06-26T13:33:10.36Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/42/e7/4bf953c6e05df90c6d351af69966384fed8e988d0e8c54dad7103b59f3ba/numpy-1.24.4-pp38-pypy38_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:95f7ac6540e95bc440ad77f56e520da5bf877f87dca58bd095288dce8940532a", size = 16705744, upload-time = "2023-06-26T13:33:36.703Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fc/dd/9106005eb477d022b60b3817ed5937a43dad8fd1f20b0610ea8a32fcb407/numpy-1.24.4-pp38-pypy38_pp73-win_amd64.whl", hash = "sha256:e98f220aa76ca2a977fe435f5b04d7b3470c0a2e6312907b37ba6068f26787f2", size = 14734290, upload-time = "2023-06-26T13:34:05.409Z" }, +] + +[[package]] +name = "numpy" +version = "2.0.2" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform == 'darwin')", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a9/75/10dd1f8116a8b796cb2c737b674e02d02e80454bda953fa7e65d8c12b016/numpy-2.0.2.tar.gz", hash = "sha256:883c987dee1880e2a864ab0dc9892292582510604156762362d9326444636e78", size = 18902015, upload-time = "2024-08-26T20:19:40.945Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/21/91/3495b3237510f79f5d81f2508f9f13fea78ebfdf07538fc7444badda173d/numpy-2.0.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:51129a29dbe56f9ca83438b706e2e69a39892b5eda6cedcb6b0c9fdc9b0d3ece", size = 21165245, upload-time = "2024-08-26T20:04:14.625Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/05/33/26178c7d437a87082d11019292dce6d3fe6f0e9026b7b2309cbf3e489b1d/numpy-2.0.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:f15975dfec0cf2239224d80e32c3170b1d168335eaedee69da84fbe9f1f9cd04", size = 13738540, upload-time = "2024-08-26T20:04:36.784Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ec/31/cc46e13bf07644efc7a4bf68df2df5fb2a1a88d0cd0da9ddc84dc0033e51/numpy-2.0.2-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:8c5713284ce4e282544c68d1c3b2c7161d38c256d2eefc93c1d683cf47683e66", size = 5300623, upload-time = "2024-08-26T20:04:46.491Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6e/16/7bfcebf27bb4f9d7ec67332ffebee4d1bf085c84246552d52dbb548600e7/numpy-2.0.2-cp310-cp310-macosx_14_0_x86_64.whl", hash = "sha256:becfae3ddd30736fe1889a37f1f580e245ba79a5855bff5f2a29cb3ccc22dd7b", size = 6901774, upload-time = "2024-08-26T20:04:58.173Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f9/a3/561c531c0e8bf082c5bef509d00d56f82e0ea7e1e3e3a7fc8fa78742a6e5/numpy-2.0.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2da5960c3cf0df7eafefd806d4e612c5e19358de82cb3c343631188991566ccd", size = 13907081, upload-time = "2024-08-26T20:05:19.098Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fa/66/f7177ab331876200ac7563a580140643d1179c8b4b6a6b0fc9838de2a9b8/numpy-2.0.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:496f71341824ed9f3d2fd36cf3ac57ae2e0165c143b55c3a035ee219413f3318", size = 19523451, upload-time = "2024-08-26T20:05:47.479Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/25/7f/0b209498009ad6453e4efc2c65bcdf0ae08a182b2b7877d7ab38a92dc542/numpy-2.0.2-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:a61ec659f68ae254e4d237816e33171497e978140353c0c2038d46e63282d0c8", size = 19927572, upload-time = "2024-08-26T20:06:17.137Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3e/df/2619393b1e1b565cd2d4c4403bdd979621e2c4dea1f8532754b2598ed63b/numpy-2.0.2-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:d731a1c6116ba289c1e9ee714b08a8ff882944d4ad631fd411106a30f083c326", size = 14400722, upload-time = "2024-08-26T20:06:39.16Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/22/ad/77e921b9f256d5da36424ffb711ae79ca3f451ff8489eeca544d0701d74a/numpy-2.0.2-cp310-cp310-win32.whl", hash = "sha256:984d96121c9f9616cd33fbd0618b7f08e0cfc9600a7ee1d6fd9b239186d19d97", size = 6472170, upload-time = "2024-08-26T20:06:50.361Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/10/05/3442317535028bc29cf0c0dd4c191a4481e8376e9f0db6bcf29703cadae6/numpy-2.0.2-cp310-cp310-win_amd64.whl", hash = "sha256:c7b0be4ef08607dd04da4092faee0b86607f111d5ae68036f16cc787e250a131", size = 15905558, upload-time = "2024-08-26T20:07:13.881Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8b/cf/034500fb83041aa0286e0fb16e7c76e5c8b67c0711bb6e9e9737a717d5fe/numpy-2.0.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:49ca4decb342d66018b01932139c0961a8f9ddc7589611158cb3c27cbcf76448", size = 21169137, upload-time = "2024-08-26T20:07:45.345Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4a/d9/32de45561811a4b87fbdee23b5797394e3d1504b4a7cf40c10199848893e/numpy-2.0.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:11a76c372d1d37437857280aa142086476136a8c0f373b2e648ab2c8f18fb195", size = 13703552, upload-time = "2024-08-26T20:08:06.666Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c1/ca/2f384720020c7b244d22508cb7ab23d95f179fcfff33c31a6eeba8d6c512/numpy-2.0.2-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:807ec44583fd708a21d4a11d94aedf2f4f3c3719035c76a2bbe1fe8e217bdc57", size = 5298957, upload-time = "2024-08-26T20:08:15.83Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0e/78/a3e4f9fb6aa4e6fdca0c5428e8ba039408514388cf62d89651aade838269/numpy-2.0.2-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:8cafab480740e22f8d833acefed5cc87ce276f4ece12fdaa2e8903db2f82897a", size = 6905573, upload-time = "2024-08-26T20:08:27.185Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a0/72/cfc3a1beb2caf4efc9d0b38a15fe34025230da27e1c08cc2eb9bfb1c7231/numpy-2.0.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a15f476a45e6e5a3a79d8a14e62161d27ad897381fecfa4a09ed5322f2085669", size = 13914330, upload-time = "2024-08-26T20:08:48.058Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ba/a8/c17acf65a931ce551fee11b72e8de63bf7e8a6f0e21add4c937c83563538/numpy-2.0.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:13e689d772146140a252c3a28501da66dfecd77490b498b168b501835041f951", size = 19534895, upload-time = "2024-08-26T20:09:16.536Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ba/86/8767f3d54f6ae0165749f84648da9dcc8cd78ab65d415494962c86fac80f/numpy-2.0.2-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:9ea91dfb7c3d1c56a0e55657c0afb38cf1eeae4544c208dc465c3c9f3a7c09f9", size = 19937253, upload-time = "2024-08-26T20:09:46.263Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/df/87/f76450e6e1c14e5bb1eae6836478b1028e096fd02e85c1c37674606ab752/numpy-2.0.2-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:c1c9307701fec8f3f7a1e6711f9089c06e6284b3afbbcd259f7791282d660a15", size = 14414074, upload-time = "2024-08-26T20:10:08.483Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5c/ca/0f0f328e1e59f73754f06e1adfb909de43726d4f24c6a3f8805f34f2b0fa/numpy-2.0.2-cp311-cp311-win32.whl", hash = "sha256:a392a68bd329eafac5817e5aefeb39038c48b671afd242710b451e76090e81f4", size = 6470640, upload-time = "2024-08-26T20:10:19.732Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/eb/57/3a3f14d3a759dcf9bf6e9eda905794726b758819df4663f217d658a58695/numpy-2.0.2-cp311-cp311-win_amd64.whl", hash = "sha256:286cd40ce2b7d652a6f22efdfc6d1edf879440e53e76a75955bc0c826c7e64dc", size = 15910230, upload-time = "2024-08-26T20:10:43.413Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/45/40/2e117be60ec50d98fa08c2f8c48e09b3edea93cfcabd5a9ff6925d54b1c2/numpy-2.0.2-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:df55d490dea7934f330006d0f81e8551ba6010a5bf035a249ef61a94f21c500b", size = 20895803, upload-time = "2024-08-26T20:11:13.916Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/46/92/1b8b8dee833f53cef3e0a3f69b2374467789e0bb7399689582314df02651/numpy-2.0.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:8df823f570d9adf0978347d1f926b2a867d5608f434a7cff7f7908c6570dcf5e", size = 13471835, upload-time = "2024-08-26T20:11:34.779Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7f/19/e2793bde475f1edaea6945be141aef6c8b4c669b90c90a300a8954d08f0a/numpy-2.0.2-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:9a92ae5c14811e390f3767053ff54eaee3bf84576d99a2456391401323f4ec2c", size = 5038499, upload-time = "2024-08-26T20:11:43.902Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e3/ff/ddf6dac2ff0dd50a7327bcdba45cb0264d0e96bb44d33324853f781a8f3c/numpy-2.0.2-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:a842d573724391493a97a62ebbb8e731f8a5dcc5d285dfc99141ca15a3302d0c", size = 6633497, upload-time = "2024-08-26T20:11:55.09Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/72/21/67f36eac8e2d2cd652a2e69595a54128297cdcb1ff3931cfc87838874bd4/numpy-2.0.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c05e238064fc0610c840d1cf6a13bf63d7e391717d247f1bf0318172e759e692", size = 13621158, upload-time = "2024-08-26T20:12:14.95Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/39/68/e9f1126d757653496dbc096cb429014347a36b228f5a991dae2c6b6cfd40/numpy-2.0.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0123ffdaa88fa4ab64835dcbde75dcdf89c453c922f18dced6e27c90d1d0ec5a", size = 19236173, upload-time = "2024-08-26T20:12:44.049Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d1/e9/1f5333281e4ebf483ba1c888b1d61ba7e78d7e910fdd8e6499667041cc35/numpy-2.0.2-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:96a55f64139912d61de9137f11bf39a55ec8faec288c75a54f93dfd39f7eb40c", size = 19634174, upload-time = "2024-08-26T20:13:13.634Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/71/af/a469674070c8d8408384e3012e064299f7a2de540738a8e414dcfd639996/numpy-2.0.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:ec9852fb39354b5a45a80bdab5ac02dd02b15f44b3804e9f00c556bf24b4bded", size = 14099701, upload-time = "2024-08-26T20:13:34.851Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d0/3d/08ea9f239d0e0e939b6ca52ad403c84a2bce1bde301a8eb4888c1c1543f1/numpy-2.0.2-cp312-cp312-win32.whl", hash = "sha256:671bec6496f83202ed2d3c8fdc486a8fc86942f2e69ff0e986140339a63bcbe5", size = 6174313, upload-time = "2024-08-26T20:13:45.653Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b2/b5/4ac39baebf1fdb2e72585c8352c56d063b6126be9fc95bd2bb5ef5770c20/numpy-2.0.2-cp312-cp312-win_amd64.whl", hash = "sha256:cfd41e13fdc257aa5778496b8caa5e856dc4896d4ccf01841daee1d96465467a", size = 15606179, upload-time = "2024-08-26T20:14:08.786Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/43/c1/41c8f6df3162b0c6ffd4437d729115704bd43363de0090c7f913cfbc2d89/numpy-2.0.2-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:9059e10581ce4093f735ed23f3b9d283b9d517ff46009ddd485f1747eb22653c", size = 21169942, upload-time = "2024-08-26T20:14:40.108Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/39/bc/fd298f308dcd232b56a4031fd6ddf11c43f9917fbc937e53762f7b5a3bb1/numpy-2.0.2-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:423e89b23490805d2a5a96fe40ec507407b8ee786d66f7328be214f9679df6dd", size = 13711512, upload-time = "2024-08-26T20:15:00.985Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/96/ff/06d1aa3eeb1c614eda245c1ba4fb88c483bee6520d361641331872ac4b82/numpy-2.0.2-cp39-cp39-macosx_14_0_arm64.whl", hash = "sha256:2b2955fa6f11907cf7a70dab0d0755159bca87755e831e47932367fc8f2f2d0b", size = 5306976, upload-time = "2024-08-26T20:15:10.876Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2d/98/121996dcfb10a6087a05e54453e28e58694a7db62c5a5a29cee14c6e047b/numpy-2.0.2-cp39-cp39-macosx_14_0_x86_64.whl", hash = "sha256:97032a27bd9d8988b9a97a8c4d2c9f2c15a81f61e2f21404d7e8ef00cb5be729", size = 6906494, upload-time = "2024-08-26T20:15:22.055Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/15/31/9dffc70da6b9bbf7968f6551967fc21156207366272c2a40b4ed6008dc9b/numpy-2.0.2-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1e795a8be3ddbac43274f18588329c72939870a16cae810c2b73461c40718ab1", size = 13912596, upload-time = "2024-08-26T20:15:42.452Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b9/14/78635daab4b07c0930c919d451b8bf8c164774e6a3413aed04a6d95758ce/numpy-2.0.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f26b258c385842546006213344c50655ff1555a9338e2e5e02a0756dc3e803dd", size = 19526099, upload-time = "2024-08-26T20:16:11.048Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/26/4c/0eeca4614003077f68bfe7aac8b7496f04221865b3a5e7cb230c9d055afd/numpy-2.0.2-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:5fec9451a7789926bcf7c2b8d187292c9f93ea30284802a0ab3f5be8ab36865d", size = 19932823, upload-time = "2024-08-26T20:16:40.171Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f1/46/ea25b98b13dccaebddf1a803f8c748680d972e00507cd9bc6dcdb5aa2ac1/numpy-2.0.2-cp39-cp39-musllinux_1_2_aarch64.whl", hash = "sha256:9189427407d88ff25ecf8f12469d4d39d35bee1db5d39fc5c168c6f088a6956d", size = 14404424, upload-time = "2024-08-26T20:17:02.604Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c8/a6/177dd88d95ecf07e722d21008b1b40e681a929eb9e329684d449c36586b2/numpy-2.0.2-cp39-cp39-win32.whl", hash = "sha256:905d16e0c60200656500c95b6b8dca5d109e23cb24abc701d41c02d74c6b3afa", size = 6476809, upload-time = "2024-08-26T20:17:13.553Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ea/2b/7fc9f4e7ae5b507c1a3a21f0f15ed03e794c1242ea8a242ac158beb56034/numpy-2.0.2-cp39-cp39-win_amd64.whl", hash = "sha256:a3f4ab0caa7f053f6797fcd4e1e25caee367db3112ef2b6ef82d749530768c73", size = 15911314, upload-time = "2024-08-26T20:17:36.72Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8f/3b/df5a870ac6a3be3a86856ce195ef42eec7ae50d2a202be1f5a4b3b340e14/numpy-2.0.2-pp39-pypy39_pp73-macosx_10_9_x86_64.whl", hash = "sha256:7f0a0c6f12e07fa94133c8a67404322845220c06a9e80e85999afe727f7438b8", size = 21025288, upload-time = "2024-08-26T20:18:07.732Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2c/97/51af92f18d6f6f2d9ad8b482a99fb74e142d71372da5d834b3a2747a446e/numpy-2.0.2-pp39-pypy39_pp73-macosx_14_0_x86_64.whl", hash = "sha256:312950fdd060354350ed123c0e25a71327d3711584beaef30cdaa93320c392d4", size = 6762793, upload-time = "2024-08-26T20:18:19.125Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/12/46/de1fbd0c1b5ccaa7f9a005b66761533e2f6a3e560096682683a223631fe9/numpy-2.0.2-pp39-pypy39_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:26df23238872200f63518dd2aa984cfca675d82469535dc7162dc2ee52d9dd5c", size = 19334885, upload-time = "2024-08-26T20:18:47.237Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/cc/dc/d330a6faefd92b446ec0f0dfea4c3207bb1fef3c4771d19cf4543efd2c78/numpy-2.0.2-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:a46288ec55ebbd58947d31d72be2c63cbf839f0a63b49cb755022310792a3385", size = 15828784, upload-time = "2024-08-26T20:19:11.19Z" }, +] + +[[package]] +name = "numpy" +version = "2.2.6" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.13.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'darwin')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'darwin')", + "(python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform == 'darwin')", + "(python_full_version == '3.10.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'darwin')", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/76/21/7d2a95e4bba9dc13d043ee156a356c0a8f0c6309dff6b21b4d71a073b8a8/numpy-2.2.6.tar.gz", hash = "sha256:e29554e2bef54a90aa5cc07da6ce955accb83f21ab5de01a62c8478897b264fd", size = 20276440, upload-time = "2025-05-17T22:38:04.611Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/9a/3e/ed6db5be21ce87955c0cbd3009f2803f59fa08df21b5df06862e2d8e2bdd/numpy-2.2.6-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:b412caa66f72040e6d268491a59f2c43bf03eb6c96dd8f0307829feb7fa2b6fb", size = 21165245, upload-time = "2025-05-17T21:27:58.555Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/22/c2/4b9221495b2a132cc9d2eb862e21d42a009f5a60e45fc44b00118c174bff/numpy-2.2.6-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:8e41fd67c52b86603a91c1a505ebaef50b3314de0213461c7a6e99c9a3beff90", size = 14360048, upload-time = "2025-05-17T21:28:21.406Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fd/77/dc2fcfc66943c6410e2bf598062f5959372735ffda175b39906d54f02349/numpy-2.2.6-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:37e990a01ae6ec7fe7fa1c26c55ecb672dd98b19c3d0e1d1f326fa13cb38d163", size = 5340542, upload-time = "2025-05-17T21:28:30.931Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7a/4f/1cb5fdc353a5f5cc7feb692db9b8ec2c3d6405453f982435efc52561df58/numpy-2.2.6-cp310-cp310-macosx_14_0_x86_64.whl", hash = "sha256:5a6429d4be8ca66d889b7cf70f536a397dc45ba6faeb5f8c5427935d9592e9cf", size = 6878301, upload-time = "2025-05-17T21:28:41.613Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/eb/17/96a3acd228cec142fcb8723bd3cc39c2a474f7dcf0a5d16731980bcafa95/numpy-2.2.6-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:efd28d4e9cd7d7a8d39074a4d44c63eda73401580c5c76acda2ce969e0a38e83", size = 14297320, upload-time = "2025-05-17T21:29:02.78Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b4/63/3de6a34ad7ad6646ac7d2f55ebc6ad439dbbf9c4370017c50cf403fb19b5/numpy-2.2.6-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fc7b73d02efb0e18c000e9ad8b83480dfcd5dfd11065997ed4c6747470ae8915", size = 16801050, upload-time = "2025-05-17T21:29:27.675Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/07/b6/89d837eddef52b3d0cec5c6ba0456c1bf1b9ef6a6672fc2b7873c3ec4e2e/numpy-2.2.6-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:74d4531beb257d2c3f4b261bfb0fc09e0f9ebb8842d82a7b4209415896adc680", size = 15807034, upload-time = "2025-05-17T21:29:51.102Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/01/c8/dc6ae86e3c61cfec1f178e5c9f7858584049b6093f843bca541f94120920/numpy-2.2.6-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:8fc377d995680230e83241d8a96def29f204b5782f371c532579b4f20607a289", size = 18614185, upload-time = "2025-05-17T21:30:18.703Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5b/c5/0064b1b7e7c89137b471ccec1fd2282fceaae0ab3a9550f2568782d80357/numpy-2.2.6-cp310-cp310-win32.whl", hash = "sha256:b093dd74e50a8cba3e873868d9e93a85b78e0daf2e98c6797566ad8044e8363d", size = 6527149, upload-time = "2025-05-17T21:30:29.788Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a3/dd/4b822569d6b96c39d1215dbae0582fd99954dcbcf0c1a13c61783feaca3f/numpy-2.2.6-cp310-cp310-win_amd64.whl", hash = "sha256:f0fd6321b839904e15c46e0d257fdd101dd7f530fe03fd6359c1ea63738703f3", size = 12904620, upload-time = "2025-05-17T21:30:48.994Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/da/a8/4f83e2aa666a9fbf56d6118faaaf5f1974d456b1823fda0a176eff722839/numpy-2.2.6-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:f9f1adb22318e121c5c69a09142811a201ef17ab257a1e66ca3025065b7f53ae", size = 21176963, upload-time = "2025-05-17T21:31:19.36Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b3/2b/64e1affc7972decb74c9e29e5649fac940514910960ba25cd9af4488b66c/numpy-2.2.6-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:c820a93b0255bc360f53eca31a0e676fd1101f673dda8da93454a12e23fc5f7a", size = 14406743, upload-time = "2025-05-17T21:31:41.087Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4a/9f/0121e375000b5e50ffdd8b25bf78d8e1a5aa4cca3f185d41265198c7b834/numpy-2.2.6-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:3d70692235e759f260c3d837193090014aebdf026dfd167834bcba43e30c2a42", size = 5352616, upload-time = "2025-05-17T21:31:50.072Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/31/0d/b48c405c91693635fbe2dcd7bc84a33a602add5f63286e024d3b6741411c/numpy-2.2.6-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:481b49095335f8eed42e39e8041327c05b0f6f4780488f61286ed3c01368d491", size = 6889579, upload-time = "2025-05-17T21:32:01.712Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/52/b8/7f0554d49b565d0171eab6e99001846882000883998e7b7d9f0d98b1f934/numpy-2.2.6-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b64d8d4d17135e00c8e346e0a738deb17e754230d7e0810ac5012750bbd85a5a", size = 14312005, upload-time = "2025-05-17T21:32:23.332Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b3/dd/2238b898e51bd6d389b7389ffb20d7f4c10066d80351187ec8e303a5a475/numpy-2.2.6-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba10f8411898fc418a521833e014a77d3ca01c15b0c6cdcce6a0d2897e6dbbdf", size = 16821570, upload-time = "2025-05-17T21:32:47.991Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/83/6c/44d0325722cf644f191042bf47eedad61c1e6df2432ed65cbe28509d404e/numpy-2.2.6-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:bd48227a919f1bafbdda0583705e547892342c26fb127219d60a5c36882609d1", size = 15818548, upload-time = "2025-05-17T21:33:11.728Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ae/9d/81e8216030ce66be25279098789b665d49ff19eef08bfa8cb96d4957f422/numpy-2.2.6-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:9551a499bf125c1d4f9e250377c1ee2eddd02e01eac6644c080162c0c51778ab", size = 18620521, upload-time = "2025-05-17T21:33:39.139Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6a/fd/e19617b9530b031db51b0926eed5345ce8ddc669bb3bc0044b23e275ebe8/numpy-2.2.6-cp311-cp311-win32.whl", hash = "sha256:0678000bb9ac1475cd454c6b8c799206af8107e310843532b04d49649c717a47", size = 6525866, upload-time = "2025-05-17T21:33:50.273Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/31/0a/f354fb7176b81747d870f7991dc763e157a934c717b67b58456bc63da3df/numpy-2.2.6-cp311-cp311-win_amd64.whl", hash = "sha256:e8213002e427c69c45a52bbd94163084025f533a55a59d6f9c5b820774ef3303", size = 12907455, upload-time = "2025-05-17T21:34:09.135Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/82/5d/c00588b6cf18e1da539b45d3598d3557084990dcc4331960c15ee776ee41/numpy-2.2.6-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:41c5a21f4a04fa86436124d388f6ed60a9343a6f767fced1a8a71c3fbca038ff", size = 20875348, upload-time = "2025-05-17T21:34:39.648Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/66/ee/560deadcdde6c2f90200450d5938f63a34b37e27ebff162810f716f6a230/numpy-2.2.6-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:de749064336d37e340f640b05f24e9e3dd678c57318c7289d222a8a2f543e90c", size = 14119362, upload-time = "2025-05-17T21:35:01.241Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3c/65/4baa99f1c53b30adf0acd9a5519078871ddde8d2339dc5a7fde80d9d87da/numpy-2.2.6-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:894b3a42502226a1cac872f840030665f33326fc3dac8e57c607905773cdcde3", size = 5084103, upload-time = "2025-05-17T21:35:10.622Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/cc/89/e5a34c071a0570cc40c9a54eb472d113eea6d002e9ae12bb3a8407fb912e/numpy-2.2.6-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:71594f7c51a18e728451bb50cc60a3ce4e6538822731b2933209a1f3614e9282", size = 6625382, upload-time = "2025-05-17T21:35:21.414Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f8/35/8c80729f1ff76b3921d5c9487c7ac3de9b2a103b1cd05e905b3090513510/numpy-2.2.6-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f2618db89be1b4e05f7a1a847a9c1c0abd63e63a1607d892dd54668dd92faf87", size = 14018462, upload-time = "2025-05-17T21:35:42.174Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8c/3d/1e1db36cfd41f895d266b103df00ca5b3cbe965184df824dec5c08c6b803/numpy-2.2.6-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fd83c01228a688733f1ded5201c678f0c53ecc1006ffbc404db9f7a899ac6249", size = 16527618, upload-time = "2025-05-17T21:36:06.711Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/61/c6/03ed30992602c85aa3cd95b9070a514f8b3c33e31124694438d88809ae36/numpy-2.2.6-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:37c0ca431f82cd5fa716eca9506aefcabc247fb27ba69c5062a6d3ade8cf8f49", size = 15505511, upload-time = "2025-05-17T21:36:29.965Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b7/25/5761d832a81df431e260719ec45de696414266613c9ee268394dd5ad8236/numpy-2.2.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fe27749d33bb772c80dcd84ae7e8df2adc920ae8297400dabec45f0dedb3f6de", size = 18313783, upload-time = "2025-05-17T21:36:56.883Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/57/0a/72d5a3527c5ebffcd47bde9162c39fae1f90138c961e5296491ce778e682/numpy-2.2.6-cp312-cp312-win32.whl", hash = "sha256:4eeaae00d789f66c7a25ac5f34b71a7035bb474e679f410e5e1a94deb24cf2d4", size = 6246506, upload-time = "2025-05-17T21:37:07.368Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/36/fa/8c9210162ca1b88529ab76b41ba02d433fd54fecaf6feb70ef9f124683f1/numpy-2.2.6-cp312-cp312-win_amd64.whl", hash = "sha256:c1f9540be57940698ed329904db803cf7a402f3fc200bfe599334c9bd84a40b2", size = 12614190, upload-time = "2025-05-17T21:37:26.213Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f9/5c/6657823f4f594f72b5471f1db1ab12e26e890bb2e41897522d134d2a3e81/numpy-2.2.6-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:0811bb762109d9708cca4d0b13c4f67146e3c3b7cf8d34018c722adb2d957c84", size = 20867828, upload-time = "2025-05-17T21:37:56.699Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/dc/9e/14520dc3dadf3c803473bd07e9b2bd1b69bc583cb2497b47000fed2fa92f/numpy-2.2.6-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:287cc3162b6f01463ccd86be154f284d0893d2b3ed7292439ea97eafa8170e0b", size = 14143006, upload-time = "2025-05-17T21:38:18.291Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4f/06/7e96c57d90bebdce9918412087fc22ca9851cceaf5567a45c1f404480e9e/numpy-2.2.6-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:f1372f041402e37e5e633e586f62aa53de2eac8d98cbfb822806ce4bbefcb74d", size = 5076765, upload-time = "2025-05-17T21:38:27.319Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/73/ed/63d920c23b4289fdac96ddbdd6132e9427790977d5457cd132f18e76eae0/numpy-2.2.6-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:55a4d33fa519660d69614a9fad433be87e5252f4b03850642f88993f7b2ca566", size = 6617736, upload-time = "2025-05-17T21:38:38.141Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/85/c5/e19c8f99d83fd377ec8c7e0cf627a8049746da54afc24ef0a0cb73d5dfb5/numpy-2.2.6-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f92729c95468a2f4f15e9bb94c432a9229d0d50de67304399627a943201baa2f", size = 14010719, upload-time = "2025-05-17T21:38:58.433Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/19/49/4df9123aafa7b539317bf6d342cb6d227e49f7a35b99c287a6109b13dd93/numpy-2.2.6-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1bc23a79bfabc5d056d106f9befb8d50c31ced2fbc70eedb8155aec74a45798f", size = 16526072, upload-time = "2025-05-17T21:39:22.638Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b2/6c/04b5f47f4f32f7c2b0e7260442a8cbcf8168b0e1a41ff1495da42f42a14f/numpy-2.2.6-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:e3143e4451880bed956e706a3220b4e5cf6172ef05fcc397f6f36a550b1dd868", size = 15503213, upload-time = "2025-05-17T21:39:45.865Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/17/0a/5cd92e352c1307640d5b6fec1b2ffb06cd0dabe7d7b8227f97933d378422/numpy-2.2.6-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:b4f13750ce79751586ae2eb824ba7e1e8dba64784086c98cdbbcc6a42112ce0d", size = 18316632, upload-time = "2025-05-17T21:40:13.331Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f0/3b/5cba2b1d88760ef86596ad0f3d484b1cbff7c115ae2429678465057c5155/numpy-2.2.6-cp313-cp313-win32.whl", hash = "sha256:5beb72339d9d4fa36522fc63802f469b13cdbe4fdab4a288f0c441b74272ebfd", size = 6244532, upload-time = "2025-05-17T21:43:46.099Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/cb/3b/d58c12eafcb298d4e6d0d40216866ab15f59e55d148a5658bb3132311fcf/numpy-2.2.6-cp313-cp313-win_amd64.whl", hash = "sha256:b0544343a702fa80c95ad5d3d608ea3599dd54d4632df855e4c8d24eb6ecfa1c", size = 12610885, upload-time = "2025-05-17T21:44:05.145Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6b/9e/4bf918b818e516322db999ac25d00c75788ddfd2d2ade4fa66f1f38097e1/numpy-2.2.6-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:0bca768cd85ae743b2affdc762d617eddf3bcf8724435498a1e80132d04879e6", size = 20963467, upload-time = "2025-05-17T21:40:44Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/61/66/d2de6b291507517ff2e438e13ff7b1e2cdbdb7cb40b3ed475377aece69f9/numpy-2.2.6-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:fc0c5673685c508a142ca65209b4e79ed6740a4ed6b2267dbba90f34b0b3cfda", size = 14225144, upload-time = "2025-05-17T21:41:05.695Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e4/25/480387655407ead912e28ba3a820bc69af9adf13bcbe40b299d454ec011f/numpy-2.2.6-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:5bd4fc3ac8926b3819797a7c0e2631eb889b4118a9898c84f585a54d475b7e40", size = 5200217, upload-time = "2025-05-17T21:41:15.903Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/aa/4a/6e313b5108f53dcbf3aca0c0f3e9c92f4c10ce57a0a721851f9785872895/numpy-2.2.6-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:fee4236c876c4e8369388054d02d0e9bb84821feb1a64dd59e137e6511a551f8", size = 6712014, upload-time = "2025-05-17T21:41:27.321Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b7/30/172c2d5c4be71fdf476e9de553443cf8e25feddbe185e0bd88b096915bcc/numpy-2.2.6-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e1dda9c7e08dc141e0247a5b8f49cf05984955246a327d4c48bda16821947b2f", size = 14077935, upload-time = "2025-05-17T21:41:49.738Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/12/fb/9e743f8d4e4d3c710902cf87af3512082ae3d43b945d5d16563f26ec251d/numpy-2.2.6-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f447e6acb680fd307f40d3da4852208af94afdfab89cf850986c3ca00562f4fa", size = 16600122, upload-time = "2025-05-17T21:42:14.046Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/12/75/ee20da0e58d3a66f204f38916757e01e33a9737d0b22373b3eb5a27358f9/numpy-2.2.6-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:389d771b1623ec92636b0786bc4ae56abafad4a4c513d36a55dce14bd9ce8571", size = 15586143, upload-time = "2025-05-17T21:42:37.464Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/76/95/bef5b37f29fc5e739947e9ce5179ad402875633308504a52d188302319c8/numpy-2.2.6-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8e9ace4a37db23421249ed236fdcdd457d671e25146786dfc96835cd951aa7c1", size = 18385260, upload-time = "2025-05-17T21:43:05.189Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/09/04/f2f83279d287407cf36a7a8053a5abe7be3622a4363337338f2585e4afda/numpy-2.2.6-cp313-cp313t-win32.whl", hash = "sha256:038613e9fb8c72b0a41f025a7e4c3f0b7a1b5d768ece4796b674c8f3fe13efff", size = 6377225, upload-time = "2025-05-17T21:43:16.254Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/67/0e/35082d13c09c02c011cf21570543d202ad929d961c02a147493cb0c2bdf5/numpy-2.2.6-cp313-cp313t-win_amd64.whl", hash = "sha256:6031dd6dfecc0cf9f668681a37648373bddd6421fff6c66ec1624eed0180ee06", size = 12771374, upload-time = "2025-05-17T21:43:35.479Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/9e/3b/d94a75f4dbf1ef5d321523ecac21ef23a3cd2ac8b78ae2aac40873590229/numpy-2.2.6-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:0b605b275d7bd0c640cad4e5d30fa701a8d59302e127e5f79138ad62762c3e3d", size = 21040391, upload-time = "2025-05-17T21:44:35.948Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/17/f4/09b2fa1b58f0fb4f7c7963a1649c64c4d315752240377ed74d9cd878f7b5/numpy-2.2.6-pp310-pypy310_pp73-macosx_14_0_x86_64.whl", hash = "sha256:7befc596a7dc9da8a337f79802ee8adb30a552a94f792b9c9d18c840055907db", size = 6786754, upload-time = "2025-05-17T21:44:47.446Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/af/30/feba75f143bdc868a1cc3f44ccfa6c4b9ec522b36458e738cd00f67b573f/numpy-2.2.6-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ce47521a4754c8f4593837384bd3424880629f718d87c5d44f8ed763edd63543", size = 16643476, upload-time = "2025-05-17T21:45:11.871Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/37/48/ac2a9584402fb6c0cd5b5d1a91dcf176b15760130dd386bbafdbfe3640bf/numpy-2.2.6-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:d042d24c90c41b54fd506da306759e06e568864df8ec17ccc17e9e884634fd00", size = 12812666, upload-time = "2025-05-17T21:45:31.426Z" }, +] + +[[package]] +name = "nvidia-cublas-cu12" +version = "12.4.5.8" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ae/71/1c91302526c45ab494c23f61c7a84aa568b8c1f9d196efa5993957faf906/nvidia_cublas_cu12-12.4.5.8-py3-none-manylinux2014_x86_64.whl", hash = "sha256:2fc8da60df463fdefa81e323eef2e36489e1c94335b5358bcb38360adf75ac9b", size = 363438805, upload-time = "2024-04-03T20:57:06.025Z" }, +] + +[[package]] +name = "nvidia-cublas-cu12" +version = "12.8.4.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/dc/61/e24b560ab2e2eaeb3c839129175fb330dfcfc29e5203196e5541a4c44682/nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:8ac4e771d5a348c551b2a426eda6193c19aa630236b418086020df5ba9667142", size = 594346921, upload-time = "2025-03-07T01:44:31.254Z" }, +] + +[[package]] +name = "nvidia-cuda-cupti-cu12" +version = "12.4.127" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/67/42/f4f60238e8194a3106d06a058d494b18e006c10bb2b915655bd9f6ea4cb1/nvidia_cuda_cupti_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:9dec60f5ac126f7bb551c055072b69d85392b13311fcc1bcda2202d172df30fb", size = 13813957, upload-time = "2024-04-03T20:55:01.564Z" }, +] + +[[package]] +name = "nvidia-cuda-cupti-cu12" +version = "12.8.90" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f8/02/2adcaa145158bf1a8295d83591d22e4103dbfd821bcaf6f3f53151ca4ffa/nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ea0cb07ebda26bb9b29ba82cda34849e73c166c18162d3913575b0c9db9a6182", size = 10248621, upload-time = "2025-03-07T01:40:21.213Z" }, +] + +[[package]] +name = "nvidia-cuda-nvrtc-cu12" +version = "12.4.127" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2c/14/91ae57cd4db3f9ef7aa99f4019cfa8d54cb4caa7e00975df6467e9725a9f/nvidia_cuda_nvrtc_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a178759ebb095827bd30ef56598ec182b85547f1508941a3d560eb7ea1fbf338", size = 24640306, upload-time = "2024-04-03T20:56:01.463Z" }, +] + +[[package]] +name = "nvidia-cuda-nvrtc-cu12" +version = "12.8.93" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/05/6b/32f747947df2da6994e999492ab306a903659555dddc0fbdeb9d71f75e52/nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:a7756528852ef889772a84c6cd89d41dfa74667e24cca16bb31f8f061e3e9994", size = 88040029, upload-time = "2025-03-07T01:42:13.562Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime-cu12" +version = "12.4.127" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ea/27/1795d86fe88ef397885f2e580ac37628ed058a92ed2c39dc8eac3adf0619/nvidia_cuda_runtime_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:64403288fa2136ee8e467cdc9c9427e0434110899d07c779f25b5c068934faa5", size = 883737, upload-time = "2024-04-03T20:54:51.355Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime-cu12" +version = "12.8.90" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0d/9b/a997b638fcd068ad6e4d53b8551a7d30fe8b404d6f1804abf1df69838932/nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adade8dcbd0edf427b7204d480d6066d33902cab2a4707dcfc48a2d0fd44ab90", size = 954765, upload-time = "2025-03-07T01:40:01.615Z" }, +] + +[[package]] +name = "nvidia-cudnn-cu12" +version = "9.1.0.70" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", +] +dependencies = [ + { name = "nvidia-cublas-cu12", version = "12.4.5.8", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9' and sys_platform != 'darwin'" }, +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/9f/fd/713452cd72343f682b1c7b9321e23829f00b842ceaedcda96e742ea0b0b3/nvidia_cudnn_cu12-9.1.0.70-py3-none-manylinux2014_x86_64.whl", hash = "sha256:165764f44ef8c61fcdfdfdbe769d687e06374059fbb388b6c89ecb0e28793a6f", size = 664752741, upload-time = "2024-04-22T15:24:15.253Z" }, +] + +[[package]] +name = "nvidia-cudnn-cu12" +version = "9.10.2.21" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", +] +dependencies = [ + { name = "nvidia-cublas-cu12", version = "12.8.4.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.9' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.9' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.9' and sys_platform != 'darwin' and sys_platform != 'linux')" }, +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ba/51/e123d997aa098c61d029f76663dedbfb9bc8dcf8c60cbd6adbe42f76d049/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:949452be657fa16687d0930933f032835951ef0892b37d2d53824d1a84dc97a8", size = 706758467, upload-time = "2025-06-06T21:54:08.597Z" }, +] + +[[package]] +name = "nvidia-cufft-cu12" +version = "11.2.1.3" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", +] +dependencies = [ + { name = "nvidia-nvjitlink-cu12", version = "12.4.127", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9' and sys_platform != 'darwin'" }, +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/27/94/3266821f65b92b3138631e9c8e7fe1fb513804ac934485a8d05776e1dd43/nvidia_cufft_cu12-11.2.1.3-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f083fc24912aa410be21fa16d157fed2055dab1cc4b6934a0e03cba69eb242b9", size = 211459117, upload-time = "2024-04-03T20:57:40.402Z" }, +] + +[[package]] +name = "nvidia-cufft-cu12" +version = "11.3.3.83" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", +] +dependencies = [ + { name = "nvidia-nvjitlink-cu12", version = "12.8.93", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.9' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.9' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.9' and sys_platform != 'darwin' and sys_platform != 'linux')" }, +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1f/13/ee4e00f30e676b66ae65b4f08cb5bcbb8392c03f54f2d5413ea99a5d1c80/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4d2dd21ec0b88cf61b62e6b43564355e5222e4a3fb394cac0db101f2dd0d4f74", size = 193118695, upload-time = "2025-03-07T01:45:27.821Z" }, +] + +[[package]] +name = "nvidia-cufile-cu12" +version = "1.13.1.3" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/bb/fe/1bcba1dfbfb8d01be8d93f07bfc502c93fa23afa6fd5ab3fc7c1df71038a/nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1d069003be650e131b21c932ec3d8969c1715379251f8d23a1860554b1cb24fc", size = 1197834, upload-time = "2025-03-07T01:45:50.723Z" }, +] + +[[package]] +name = "nvidia-curand-cu12" +version = "10.3.5.147" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8a/6d/44ad094874c6f1b9c654f8ed939590bdc408349f137f9b98a3a23ccec411/nvidia_curand_cu12-10.3.5.147-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a88f583d4e0bb643c49743469964103aa59f7f708d862c3ddb0fc07f851e3b8b", size = 56305206, upload-time = "2024-04-03T20:58:08.722Z" }, +] + +[[package]] +name = "nvidia-curand-cu12" +version = "10.3.9.90" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fb/aa/6584b56dc84ebe9cf93226a5cde4d99080c8e90ab40f0c27bda7a0f29aa1/nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:b32331d4f4df5d6eefa0554c565b626c7216f87a06a4f56fab27c3b68a830ec9", size = 63619976, upload-time = "2025-03-07T01:46:23.323Z" }, +] + +[[package]] +name = "nvidia-cusolver-cu12" +version = "11.6.1.9" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", +] +dependencies = [ + { name = "nvidia-cublas-cu12", version = "12.4.5.8", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9' and sys_platform != 'darwin'" }, + { name = "nvidia-cusparse-cu12", version = "12.3.1.170", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9' and sys_platform != 'darwin'" }, + { name = "nvidia-nvjitlink-cu12", version = "12.4.127", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9' and sys_platform != 'darwin'" }, +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3a/e1/5b9089a4b2a4790dfdea8b3a006052cfecff58139d5a4e34cb1a51df8d6f/nvidia_cusolver_cu12-11.6.1.9-py3-none-manylinux2014_x86_64.whl", hash = "sha256:19e33fa442bcfd085b3086c4ebf7e8debc07cfe01e11513cc6d332fd918ac260", size = 127936057, upload-time = "2024-04-03T20:58:28.735Z" }, +] + +[[package]] +name = "nvidia-cusolver-cu12" +version = "11.7.3.90" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", +] +dependencies = [ + { name = "nvidia-cublas-cu12", version = "12.8.4.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.9' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.9' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.9' and sys_platform != 'darwin' and sys_platform != 'linux')" }, + { name = "nvidia-cusparse-cu12", version = "12.5.8.93", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.9' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.9' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.9' and sys_platform != 'darwin' and sys_platform != 'linux')" }, + { name = "nvidia-nvjitlink-cu12", version = "12.8.93", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.9' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.9' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.9' and sys_platform != 'darwin' and sys_platform != 'linux')" }, +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/85/48/9a13d2975803e8cf2777d5ed57b87a0b6ca2cc795f9a4f59796a910bfb80/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:4376c11ad263152bd50ea295c05370360776f8c3427b30991df774f9fb26c450", size = 267506905, upload-time = "2025-03-07T01:47:16.273Z" }, +] + +[[package]] +name = "nvidia-cusparse-cu12" +version = "12.3.1.170" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", +] +dependencies = [ + { name = "nvidia-nvjitlink-cu12", version = "12.4.127", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9' and sys_platform != 'darwin'" }, +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/db/f7/97a9ea26ed4bbbfc2d470994b8b4f338ef663be97b8f677519ac195e113d/nvidia_cusparse_cu12-12.3.1.170-py3-none-manylinux2014_x86_64.whl", hash = "sha256:ea4f11a2904e2a8dc4b1833cc1b5181cde564edd0d5cd33e3c168eff2d1863f1", size = 207454763, upload-time = "2024-04-03T20:58:59.995Z" }, +] + +[[package]] +name = "nvidia-cusparse-cu12" +version = "12.5.8.93" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", +] +dependencies = [ + { name = "nvidia-nvjitlink-cu12", version = "12.8.93", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.9' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.9' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.9' and sys_platform != 'darwin' and sys_platform != 'linux')" }, +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c2/f5/e1854cb2f2bcd4280c44736c93550cc300ff4b8c95ebe370d0aa7d2b473d/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ec05d76bbbd8b61b06a80e1eaf8cf4959c3d4ce8e711b65ebd0443bb0ebb13b", size = 288216466, upload-time = "2025-03-07T01:48:13.779Z" }, +] + +[[package]] +name = "nvidia-cusparselt-cu12" +version = "0.7.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/56/79/12978b96bd44274fe38b5dde5cfb660b1d114f70a65ef962bcbbed99b549/nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f1bb701d6b930d5a7cea44c19ceb973311500847f81b634d802b7b539dc55623", size = 287193691, upload-time = "2025-02-26T00:15:44.104Z" }, +] + +[[package]] +name = "nvidia-nccl-cu12" +version = "2.21.5" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/df/99/12cd266d6233f47d00daf3a72739872bdc10267d0383508b0b9c84a18bb6/nvidia_nccl_cu12-2.21.5-py3-none-manylinux2014_x86_64.whl", hash = "sha256:8579076d30a8c24988834445f8d633c697d42397e92ffc3f63fa26766d25e0a0", size = 188654414, upload-time = "2024-04-03T15:32:57.427Z" }, +] + +[[package]] +name = "nvidia-nccl-cu12" +version = "2.27.3" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5c/5b/4e4fff7bad39adf89f735f2bc87248c81db71205b62bcc0d5ca5b606b3c3/nvidia_nccl_cu12-2.27.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adf27ccf4238253e0b826bce3ff5fa532d65fc42322c8bfdfaf28024c0fbe039", size = 322364134, upload-time = "2025-06-03T21:58:04.013Z" }, +] + +[[package]] +name = "nvidia-nvjitlink-cu12" +version = "12.4.127" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ff/ff/847841bacfbefc97a00036e0fce5a0f086b640756dc38caea5e1bb002655/nvidia_nvjitlink_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:06b3b9b25bf3f8af351d664978ca26a16d2c5127dbd53c0497e28d1fb9611d57", size = 21066810, upload-time = "2024-04-03T20:59:46.957Z" }, +] + +[[package]] +name = "nvidia-nvjitlink-cu12" +version = "12.8.93" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f6/74/86a07f1d0f42998ca31312f998bd3b9a7eff7f52378f4f270c8679c77fb9/nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:81ff63371a7ebd6e6451970684f916be2eab07321b73c9d244dc2b4da7f73b88", size = 39254836, upload-time = "2025-03-07T01:49:55.661Z" }, +] + +[[package]] +name = "nvidia-nvtx-cu12" +version = "12.4.127" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/87/20/199b8713428322a2f22b722c62b8cc278cc53dffa9705d744484b5035ee9/nvidia_nvtx_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:781e950d9b9f60d8241ccea575b32f5105a5baf4c2351cab5256a24869f12a1a", size = 99144, upload-time = "2024-04-03T20:56:12.406Z" }, +] + +[[package]] +name = "nvidia-nvtx-cu12" +version = "12.8.90" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a2/eb/86626c1bbc2edb86323022371c39aa48df6fd8b0a1647bc274577f72e90b/nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b17e2001cc0d751a5bc2c6ec6d26ad95913324a4adb86788c944f8ce9ba441f", size = 89954, upload-time = "2025-03-07T01:42:44.131Z" }, +] + +[[package]] +name = "omegaconf" +version = "2.3.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +dependencies = [ + { name = "antlr4-python3-runtime" }, + { name = "pyyaml" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/09/48/6388f1bb9da707110532cb70ec4d2822858ddfb44f1cdf1233c20a80ea4b/omegaconf-2.3.0.tar.gz", hash = "sha256:d5d4b6d29955cc50ad50c46dc269bcd92c6e00f5f90d23ab5fee7bfca4ba4cc7", size = 3298120, upload-time = "2022-12-08T20:59:22.753Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e3/94/1843518e420fa3ed6919835845df698c7e27e183cb997394e4a670973a65/omegaconf-2.3.0-py3-none-any.whl", hash = "sha256:7b4df175cdb08ba400f45cae3bdcae7ba8365db4d165fc65fd04b050ab63b46b", size = 79500, upload-time = "2022-12-08T20:59:19.686Z" }, +] + +[[package]] +name = "openpyxl" +version = "3.1.5" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +dependencies = [ + { name = "et-xmlfile" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3d/f9/88d94a75de065ea32619465d2f77b29a0469500e99012523b91cc4141cd1/openpyxl-3.1.5.tar.gz", hash = "sha256:cf0e3cf56142039133628b5acffe8ef0c12bc902d2aadd3e0fe5878dc08d1050", size = 186464, upload-time = "2024-06-28T14:03:44.161Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c0/da/977ded879c29cbd04de313843e76868e6e13408a94ed6b987245dc7c8506/openpyxl-3.1.5-py2.py3-none-any.whl", hash = "sha256:5282c12b107bffeef825f4617dc029afaf41d0ea60823bbb665ef3079dc79de2", size = 250910, upload-time = "2024-06-28T14:03:41.161Z" }, +] + +[[package]] +name = "packaging" +version = "25.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a1/d4/1fc4078c65507b51b96ca8f8c3ba19e6a61c8253c72794544580a7b6c24d/packaging-25.0.tar.gz", hash = "sha256:d443872c98d677bf60f6a1f2f8c1cb748e8fe762d2bf9d3148b5599295b0fc4f", size = 165727, upload-time = "2025-04-19T11:48:59.673Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/20/12/38679034af332785aac8774540895e234f4d07f7545804097de4b666afd8/packaging-25.0-py3-none-any.whl", hash = "sha256:29572ef2b1f17581046b3a2227d5c611fb25ec70ca1ba8554b24b0e69331a484", size = 66469, upload-time = "2025-04-19T11:48:57.875Z" }, +] + +[[package]] +name = "panda" +version = "0.3.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +dependencies = [ + { name = "requests", version = "2.32.4", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "requests", version = "2.32.5", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9'" }, + { name = "setuptools", version = "75.3.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "setuptools", version = "80.9.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/79/03/74996420528fe488ce17c42b6400531c8067d7eb661c304fa3aa8fdad17c/panda-0.3.1.tar.gz", hash = "sha256:f213b848f09268b3e9fce0e103155ab003217c0e27f6048b6194e7f90bb2b716", size = 5833, upload-time = "2015-08-03T13:33:16.589Z" } + +[[package]] +name = "platformdirs" +version = "4.3.6" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/13/fc/128cc9cb8f03208bdbf93d3aa862e16d376844a14f9a0ce5cf4507372de4/platformdirs-4.3.6.tar.gz", hash = "sha256:357fb2acbc885b0419afd3ce3ed34564c13c9b95c89360cd9563f73aa5e2b907", size = 21302, upload-time = "2024-09-17T19:06:50.688Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3c/a6/bc1012356d8ece4d66dd75c4b9fc6c1f6650ddd5991e421177d9f8f671be/platformdirs-4.3.6-py3-none-any.whl", hash = "sha256:73e575e1408ab8103900836b97580d5307456908a03e92031bab39e4554cc3fb", size = 18439, upload-time = "2024-09-17T19:06:49.212Z" }, +] + +[[package]] +name = "platformdirs" +version = "4.4.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.13.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'darwin')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'darwin')", + "(python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform == 'darwin')", + "(python_full_version == '3.10.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'darwin')", + "(python_full_version == '3.9.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform == 'darwin')", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/23/e8/21db9c9987b0e728855bd57bff6984f67952bea55d6f75e055c46b5383e8/platformdirs-4.4.0.tar.gz", hash = "sha256:ca753cf4d81dc309bc67b0ea38fd15dc97bc30ce419a7f58d13eb3bf14c4febf", size = 21634, upload-time = "2025-08-26T14:32:04.268Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/40/4b/2028861e724d3bd36227adfa20d3fd24c3fc6d52032f4a93c133be5d17ce/platformdirs-4.4.0-py3-none-any.whl", hash = "sha256:abd01743f24e5287cd7a5db3752faf1a2d65353f38ec26d98e25a6db65958c85", size = 18654, upload-time = "2025-08-26T14:32:02.735Z" }, +] + +[[package]] +name = "pooch" +version = "1.8.2" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +dependencies = [ + { name = "packaging" }, + { name = "platformdirs", version = "4.3.6", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "platformdirs", version = "4.4.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9'" }, + { name = "requests", version = "2.32.4", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "requests", version = "2.32.5", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c6/77/b3d3e00c696c16cf99af81ef7b1f5fe73bd2a307abca41bd7605429fe6e5/pooch-1.8.2.tar.gz", hash = "sha256:76561f0de68a01da4df6af38e9955c4c9d1a5c90da73f7e40276a5728ec83d10", size = 59353, upload-time = "2024-06-06T16:53:46.224Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a8/87/77cc11c7a9ea9fd05503def69e3d18605852cd0d4b0d3b8f15bbeb3ef1d1/pooch-1.8.2-py3-none-any.whl", hash = "sha256:3529a57096f7198778a5ceefd5ac3ef0e4d06a6ddaf9fc2d609b806f25302c47", size = 64574, upload-time = "2024-06-06T16:53:44.343Z" }, +] + +[[package]] +name = "protobuf" +version = "3.20.3" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/55/5b/e3d951e34f8356e5feecacd12a8e3b258a1da6d9a03ad1770f28925f29bc/protobuf-3.20.3.tar.gz", hash = "sha256:2e3427429c9cffebf259491be0af70189607f365c2f41c7c3764af6f337105f2", size = 216768, upload-time = "2022-09-29T22:39:47.592Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/28/55/b80e8567ec327c060fa39b242392e25690c8899c489ecd7bb65b46b7bb55/protobuf-3.20.3-cp310-cp310-manylinux2014_aarch64.whl", hash = "sha256:f4bd856d702e5b0d96a00ec6b307b0f51c1982c2bf9c0052cf9019e9a544ba99", size = 918427, upload-time = "2022-09-29T22:10:52.591Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/31/be/80a9c6f16dfa4d41be3edbe655349778ae30882407fa8275eb46b4d34854/protobuf-3.20.3-cp310-cp310-manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:9aae4406ea63d825636cc11ffb34ad3379335803216ee3a856787bcf5ccc751e", size = 1051042, upload-time = "2022-09-29T22:10:54.629Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/db/96/948d3fcc1fa816e7ae1d27af59b9d8c5c5e582f3994fd14394f31da95b99/protobuf-3.20.3-cp310-cp310-win32.whl", hash = "sha256:28545383d61f55b57cf4df63eebd9827754fd2dc25f80c5253f9184235db242c", size = 780167, upload-time = "2022-09-29T22:10:56.276Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6f/5e/fc6feb366b0a9f28e0a2de3b062667c521cd9517d4ff55077b8f351ba2f3/protobuf-3.20.3-cp310-cp310-win_amd64.whl", hash = "sha256:67a3598f0a2dcbc58d02dd1928544e7d88f764b47d4a286202913f0b2801c2e7", size = 904029, upload-time = "2022-09-29T22:10:57.866Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/9f/1a/6848ed1669a6c70bf947d25d64ce6dcc65ccec06e917072df516944fa17e/protobuf-3.20.3-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:de78575669dddf6099a8a0f46a27e82a1783c557ccc38ee620ed8cc96d3be7d7", size = 982917, upload-time = "2022-09-29T22:11:12.347Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b5/b6/ec87636b9381137f17292851461902ceac7632a00476c2afbcd864ed5447/protobuf-3.20.3-cp38-cp38-manylinux2014_aarch64.whl", hash = "sha256:f4c42102bc82a51108e449cbb32b19b180022941c727bac0cfd50170341f16ee", size = 918639, upload-time = "2022-09-29T22:11:14.326Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/da/e4/4d62585593e9f962cb02614534f62f930de6a80a0a3784282094a01919b2/protobuf-3.20.3-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.whl", hash = "sha256:44246bab5dd4b7fbd3c0c80b6f16686808fab0e4aca819ade6e8d294a29c7050", size = 1019323, upload-time = "2022-09-29T22:11:16.254Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3c/14/16ef7da61d30a519d84e6841d096fe446c4d8a2c39b083b0376b4785f1f3/protobuf-3.20.3-cp38-cp38-win32.whl", hash = "sha256:c02ce36ec760252242a33967d51c289fd0e1c0e6e5cc9397e2279177716add86", size = 780325, upload-time = "2022-09-29T22:11:17.83Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/32/f8/52f598bceb16fe365f4ef8e957ac8890aeb56abf97d365ff5abd8c1250cf/protobuf-3.20.3-cp38-cp38-win_amd64.whl", hash = "sha256:447d43819997825d4e71bf5769d869b968ce96848b6479397e29fc24c4a5dfe9", size = 904419, upload-time = "2022-09-29T22:11:19.324Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/00/e7/d23c439c55c90ae2e52184363162f7079ca3e7d86205b411d4e9dc266f81/protobuf-3.20.3-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:398a9e0c3eaceb34ec1aee71894ca3299605fa8e761544934378bbc6c97de23b", size = 982826, upload-time = "2022-09-29T22:11:20.978Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/99/25/5825472ecd911f4ac2ac4e9ab039a48b6d03874e2add92fb633e080bf3eb/protobuf-3.20.3-cp39-cp39-manylinux2014_aarch64.whl", hash = "sha256:bf01b5720be110540be4286e791db73f84a2b721072a3711efff6c324cdf074b", size = 918423, upload-time = "2022-09-29T22:11:22.651Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c7/df/ec3ecb8c940b36121c7b77c10acebf3d1c736498aa2f1fe3b6231ee44e76/protobuf-3.20.3-cp39-cp39-manylinux_2_5_x86_64.manylinux1_x86_64.whl", hash = "sha256:daa564862dd0d39c00f8086f88700fdbe8bc717e993a21e90711acfed02f2402", size = 1019250, upload-time = "2022-09-29T22:11:24.259Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/36/8b/433071fed0058322090a55021bdc8da76d16c7bc9823f5795797803dd6d0/protobuf-3.20.3-cp39-cp39-win32.whl", hash = "sha256:819559cafa1a373b7096a482b504ae8a857c89593cf3a25af743ac9ecbd23480", size = 780270, upload-time = "2022-09-29T22:11:25.927Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/11/a5/e52b731415ad6ef3d841e9e6e337a690249e800cc7c06f0749afab26348c/protobuf-3.20.3-cp39-cp39-win_amd64.whl", hash = "sha256:03038ac1cfbc41aa21f6afcbcd357281d7521b4157926f30ebecc8d4ea59dcb7", size = 904215, upload-time = "2022-09-29T22:11:28.056Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8d/14/619e24a4c70df2901e1f4dbc50a6291eb63a759172558df326347dce1f0d/protobuf-3.20.3-py2.py3-none-any.whl", hash = "sha256:a7ca6d488aa8ff7f329d4c545b2dbad8ac31464f1d8b1c87ad1346717731e4db", size = 162128, upload-time = "2022-09-29T22:39:44.547Z" }, +] + +[[package]] +name = "psutil" +version = "7.0.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2a/80/336820c1ad9286a4ded7e845b2eccfcb27851ab8ac6abece774a6ff4d3de/psutil-7.0.0.tar.gz", hash = "sha256:7be9c3eba38beccb6495ea33afd982a44074b78f28c434a1f51cc07fd315c456", size = 497003, upload-time = "2025-02-13T21:54:07.946Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ed/e6/2d26234410f8b8abdbf891c9da62bee396583f713fb9f3325a4760875d22/psutil-7.0.0-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:101d71dc322e3cffd7cea0650b09b3d08b8e7c4109dd6809fe452dfd00e58b25", size = 238051, upload-time = "2025-02-13T21:54:12.36Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/04/8b/30f930733afe425e3cbfc0e1468a30a18942350c1a8816acfade80c005c4/psutil-7.0.0-cp36-abi3-macosx_11_0_arm64.whl", hash = "sha256:39db632f6bb862eeccf56660871433e111b6ea58f2caea825571951d4b6aa3da", size = 239535, upload-time = "2025-02-13T21:54:16.07Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2a/ed/d362e84620dd22876b55389248e522338ed1bf134a5edd3b8231d7207f6d/psutil-7.0.0-cp36-abi3-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1fcee592b4c6f146991ca55919ea3d1f8926497a713ed7faaf8225e174581e91", size = 275004, upload-time = "2025-02-13T21:54:18.662Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/bf/b9/b0eb3f3cbcb734d930fdf839431606844a825b23eaf9a6ab371edac8162c/psutil-7.0.0-cp36-abi3-manylinux_2_12_x86_64.manylinux2010_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4b1388a4f6875d7e2aff5c4ca1cc16c545ed41dd8bb596cefea80111db353a34", size = 277986, upload-time = "2025-02-13T21:54:21.811Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/eb/a2/709e0fe2f093556c17fbafda93ac032257242cabcc7ff3369e2cb76a97aa/psutil-7.0.0-cp36-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a5f098451abc2828f7dc6b58d44b532b22f2088f4999a937557b603ce72b1993", size = 279544, upload-time = "2025-02-13T21:54:24.68Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/50/e6/eecf58810b9d12e6427369784efe814a1eec0f492084ce8eb8f4d89d6d61/psutil-7.0.0-cp37-abi3-win32.whl", hash = "sha256:ba3fcef7523064a6c9da440fc4d6bd07da93ac726b5733c29027d7dc95b39d99", size = 241053, upload-time = "2025-02-13T21:54:34.31Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/50/1b/6921afe68c74868b4c9fa424dad3be35b095e16687989ebbb50ce4fceb7c/psutil-7.0.0-cp37-abi3-win_amd64.whl", hash = "sha256:4cf3d4eb1aa9b348dec30105c55cd9b7d4629285735a102beb4441e38db90553", size = 244885, upload-time = "2025-02-13T21:54:37.486Z" }, +] + +[[package]] +name = "pycparser" +version = "2.22" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1d/b2/31537cf4b1ca988837256c910a668b553fceb8f069bedc4b1c826024b52c/pycparser-2.22.tar.gz", hash = "sha256:491c8be9c040f5390f5bf44a5b07752bd07f56edf992381b05c701439eec10f6", size = 172736, upload-time = "2024-03-30T13:22:22.564Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/13/a3/a812df4e2dd5696d1f351d58b8fe16a405b234ad2886a0dab9183fb78109/pycparser-2.22-py3-none-any.whl", hash = "sha256:c3702b6d3dd8c7abc1afa565d7e63d53a1d0bd86cdc24edd75470f4de499cfcc", size = 117552, upload-time = "2024-03-30T13:22:20.476Z" }, +] + +[[package]] +name = "pynndescent" +version = "0.5.13" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +dependencies = [ + { name = "joblib", version = "1.4.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "joblib", version = "1.5.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9'" }, + { name = "llvmlite", version = "0.41.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "llvmlite", version = "0.43.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, + { name = "llvmlite", version = "0.44.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.10'" }, + { name = "numba", version = "0.58.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "numba", version = "0.60.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, + { name = "numba", version = "0.61.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.10'" }, + { name = "scikit-learn", version = "1.3.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "scikit-learn", version = "1.6.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, + { name = "scikit-learn", version = "1.7.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.10'" }, + { name = "scipy", version = "1.10.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "scipy", version = "1.13.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, + { name = "scipy", version = "1.15.3", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.10.*'" }, + { name = "scipy", version = "1.16.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.11'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7e/58/560a4db5eb3794d922fe55804b10326534ded3d971e1933c1eef91193f5e/pynndescent-0.5.13.tar.gz", hash = "sha256:d74254c0ee0a1eeec84597d5fe89fedcf778593eeabe32c2f97412934a9800fb", size = 2975955, upload-time = "2024-06-17T15:48:32.914Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d2/53/d23a97e0a2c690d40b165d1062e2c4ccc796be458a1ce59f6ba030434663/pynndescent-0.5.13-py3-none-any.whl", hash = "sha256:69aabb8f394bc631b6ac475a1c7f3994c54adf3f51cd63b2730fefba5771b949", size = 56850, upload-time = "2024-06-17T15:48:31.184Z" }, +] + +[[package]] +name = "pyyaml" +version = "6.0.2" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/54/ed/79a089b6be93607fa5cdaedf301d7dfb23af5f25c398d5ead2525b063e17/pyyaml-6.0.2.tar.gz", hash = "sha256:d584d9ec91ad65861cc08d42e834324ef890a082e591037abe114850ff7bbc3e", size = 130631, upload-time = "2024-08-06T20:33:50.674Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/9b/95/a3fac87cb7158e231b5a6012e438c647e1a87f09f8e0d123acec8ab8bf71/PyYAML-6.0.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:0a9a2848a5b7feac301353437eb7d5957887edbf81d56e903999a75a3d743086", size = 184199, upload-time = "2024-08-06T20:31:40.178Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c7/7a/68bd47624dab8fd4afbfd3c48e3b79efe09098ae941de5b58abcbadff5cb/PyYAML-6.0.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:29717114e51c84ddfba879543fb232a6ed60086602313ca38cce623c1d62cfbf", size = 171758, upload-time = "2024-08-06T20:31:42.173Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/49/ee/14c54df452143b9ee9f0f29074d7ca5516a36edb0b4cc40c3f280131656f/PyYAML-6.0.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8824b5a04a04a047e72eea5cec3bc266db09e35de6bdfe34c9436ac5ee27d237", size = 718463, upload-time = "2024-08-06T20:31:44.263Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4d/61/de363a97476e766574650d742205be468921a7b532aa2499fcd886b62530/PyYAML-6.0.2-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:7c36280e6fb8385e520936c3cb3b8042851904eba0e58d277dca80a5cfed590b", size = 719280, upload-time = "2024-08-06T20:31:50.199Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6b/4e/1523cb902fd98355e2e9ea5e5eb237cbc5f3ad5f3075fa65087aa0ecb669/PyYAML-6.0.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ec031d5d2feb36d1d1a24380e4db6d43695f3748343d99434e6f5f9156aaa2ed", size = 751239, upload-time = "2024-08-06T20:31:52.292Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b7/33/5504b3a9a4464893c32f118a9cc045190a91637b119a9c881da1cf6b7a72/PyYAML-6.0.2-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:936d68689298c36b53b29f23c6dbb74de12b4ac12ca6cfe0e047bedceea56180", size = 695802, upload-time = "2024-08-06T20:31:53.836Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5c/20/8347dcabd41ef3a3cdc4f7b7a2aff3d06598c8779faa189cdbf878b626a4/PyYAML-6.0.2-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:23502f431948090f597378482b4812b0caae32c22213aecf3b55325e049a6c68", size = 720527, upload-time = "2024-08-06T20:31:55.565Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/be/aa/5afe99233fb360d0ff37377145a949ae258aaab831bde4792b32650a4378/PyYAML-6.0.2-cp310-cp310-win32.whl", hash = "sha256:2e99c6826ffa974fe6e27cdb5ed0021786b03fc98e5ee3c5bfe1fd5015f42b99", size = 144052, upload-time = "2024-08-06T20:31:56.914Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b5/84/0fa4b06f6d6c958d207620fc60005e241ecedceee58931bb20138e1e5776/PyYAML-6.0.2-cp310-cp310-win_amd64.whl", hash = "sha256:a4d3091415f010369ae4ed1fc6b79def9416358877534caf6a0fdd2146c87a3e", size = 161774, upload-time = "2024-08-06T20:31:58.304Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f8/aa/7af4e81f7acba21a4c6be026da38fd2b872ca46226673c89a758ebdc4fd2/PyYAML-6.0.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:cc1c1159b3d456576af7a3e4d1ba7e6924cb39de8f67111c735f6fc832082774", size = 184612, upload-time = "2024-08-06T20:32:03.408Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8b/62/b9faa998fd185f65c1371643678e4d58254add437edb764a08c5a98fb986/PyYAML-6.0.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:1e2120ef853f59c7419231f3bf4e7021f1b936f6ebd222406c3b60212205d2ee", size = 172040, upload-time = "2024-08-06T20:32:04.926Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ad/0c/c804f5f922a9a6563bab712d8dcc70251e8af811fce4524d57c2c0fd49a4/PyYAML-6.0.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5d225db5a45f21e78dd9358e58a98702a0302f2659a3c6cd320564b75b86f47c", size = 736829, upload-time = "2024-08-06T20:32:06.459Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/51/16/6af8d6a6b210c8e54f1406a6b9481febf9c64a3109c541567e35a49aa2e7/PyYAML-6.0.2-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:5ac9328ec4831237bec75defaf839f7d4564be1e6b25ac710bd1a96321cc8317", size = 764167, upload-time = "2024-08-06T20:32:08.338Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/75/e4/2c27590dfc9992f73aabbeb9241ae20220bd9452df27483b6e56d3975cc5/PyYAML-6.0.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3ad2a3decf9aaba3d29c8f537ac4b243e36bef957511b4766cb0057d32b0be85", size = 762952, upload-time = "2024-08-06T20:32:14.124Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/9b/97/ecc1abf4a823f5ac61941a9c00fe501b02ac3ab0e373c3857f7d4b83e2b6/PyYAML-6.0.2-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:ff3824dc5261f50c9b0dfb3be22b4567a6f938ccce4587b38952d85fd9e9afe4", size = 735301, upload-time = "2024-08-06T20:32:16.17Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/45/73/0f49dacd6e82c9430e46f4a027baa4ca205e8b0a9dce1397f44edc23559d/PyYAML-6.0.2-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:797b4f722ffa07cc8d62053e4cff1486fa6dc094105d13fea7b1de7d8bf71c9e", size = 756638, upload-time = "2024-08-06T20:32:18.555Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/22/5f/956f0f9fc65223a58fbc14459bf34b4cc48dec52e00535c79b8db361aabd/PyYAML-6.0.2-cp311-cp311-win32.whl", hash = "sha256:11d8f3dd2b9c1207dcaf2ee0bbbfd5991f571186ec9cc78427ba5bd32afae4b5", size = 143850, upload-time = "2024-08-06T20:32:19.889Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ed/23/8da0bbe2ab9dcdd11f4f4557ccaf95c10b9811b13ecced089d43ce59c3c8/PyYAML-6.0.2-cp311-cp311-win_amd64.whl", hash = "sha256:e10ce637b18caea04431ce14fabcf5c64a1c61ec9c56b071a4b7ca131ca52d44", size = 161980, upload-time = "2024-08-06T20:32:21.273Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/86/0c/c581167fc46d6d6d7ddcfb8c843a4de25bdd27e4466938109ca68492292c/PyYAML-6.0.2-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:c70c95198c015b85feafc136515252a261a84561b7b1d51e3384e0655ddf25ab", size = 183873, upload-time = "2024-08-06T20:32:25.131Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a8/0c/38374f5bb272c051e2a69281d71cba6fdb983413e6758b84482905e29a5d/PyYAML-6.0.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ce826d6ef20b1bc864f0a68340c8b3287705cae2f8b4b1d932177dcc76721725", size = 173302, upload-time = "2024-08-06T20:32:26.511Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c3/93/9916574aa8c00aa06bbac729972eb1071d002b8e158bd0e83a3b9a20a1f7/PyYAML-6.0.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1f71ea527786de97d1a0cc0eacd1defc0985dcf6b3f17bb77dcfc8c34bec4dc5", size = 739154, upload-time = "2024-08-06T20:32:28.363Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/95/0f/b8938f1cbd09739c6da569d172531567dbcc9789e0029aa070856f123984/PyYAML-6.0.2-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:9b22676e8097e9e22e36d6b7bda33190d0d400f345f23d4065d48f4ca7ae0425", size = 766223, upload-time = "2024-08-06T20:32:30.058Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b9/2b/614b4752f2e127db5cc206abc23a8c19678e92b23c3db30fc86ab731d3bd/PyYAML-6.0.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:80bab7bfc629882493af4aa31a4cfa43a4c57c83813253626916b8c7ada83476", size = 767542, upload-time = "2024-08-06T20:32:31.881Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d4/00/dd137d5bcc7efea1836d6264f049359861cf548469d18da90cd8216cf05f/PyYAML-6.0.2-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:0833f8694549e586547b576dcfaba4a6b55b9e96098b36cdc7ebefe667dfed48", size = 731164, upload-time = "2024-08-06T20:32:37.083Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c9/1f/4f998c900485e5c0ef43838363ba4a9723ac0ad73a9dc42068b12aaba4e4/PyYAML-6.0.2-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8b9c7197f7cb2738065c481a0461e50ad02f18c78cd75775628afb4d7137fb3b", size = 756611, upload-time = "2024-08-06T20:32:38.898Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/df/d1/f5a275fdb252768b7a11ec63585bc38d0e87c9e05668a139fea92b80634c/PyYAML-6.0.2-cp312-cp312-win32.whl", hash = "sha256:ef6107725bd54b262d6dedcc2af448a266975032bc85ef0172c5f059da6325b4", size = 140591, upload-time = "2024-08-06T20:32:40.241Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0c/e8/4f648c598b17c3d06e8753d7d13d57542b30d56e6c2dedf9c331ae56312e/PyYAML-6.0.2-cp312-cp312-win_amd64.whl", hash = "sha256:7e7401d0de89a9a855c839bc697c079a4af81cf878373abd7dc625847d25cbd8", size = 156338, upload-time = "2024-08-06T20:32:41.93Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ef/e3/3af305b830494fa85d95f6d95ef7fa73f2ee1cc8ef5b495c7c3269fb835f/PyYAML-6.0.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:efdca5630322a10774e8e98e1af481aad470dd62c3170801852d752aa7a783ba", size = 181309, upload-time = "2024-08-06T20:32:43.4Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/45/9f/3b1c20a0b7a3200524eb0076cc027a970d320bd3a6592873c85c92a08731/PyYAML-6.0.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:50187695423ffe49e2deacb8cd10510bc361faac997de9efef88badc3bb9e2d1", size = 171679, upload-time = "2024-08-06T20:32:44.801Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7c/9a/337322f27005c33bcb656c655fa78325b730324c78620e8328ae28b64d0c/PyYAML-6.0.2-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0ffe8360bab4910ef1b9e87fb812d8bc0a308b0d0eef8c8f44e0254ab3b07133", size = 733428, upload-time = "2024-08-06T20:32:46.432Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a3/69/864fbe19e6c18ea3cc196cbe5d392175b4cf3d5d0ac1403ec3f2d237ebb5/PyYAML-6.0.2-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:17e311b6c678207928d649faa7cb0d7b4c26a0ba73d41e99c4fff6b6c3276484", size = 763361, upload-time = "2024-08-06T20:32:51.188Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/04/24/b7721e4845c2f162d26f50521b825fb061bc0a5afcf9a386840f23ea19fa/PyYAML-6.0.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:70b189594dbe54f75ab3a1acec5f1e3faa7e8cf2f1e08d9b561cb41b845f69d5", size = 759523, upload-time = "2024-08-06T20:32:53.019Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2b/b2/e3234f59ba06559c6ff63c4e10baea10e5e7df868092bf9ab40e5b9c56b6/PyYAML-6.0.2-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:41e4e3953a79407c794916fa277a82531dd93aad34e29c2a514c2c0c5fe971cc", size = 726660, upload-time = "2024-08-06T20:32:54.708Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fe/0f/25911a9f080464c59fab9027482f822b86bf0608957a5fcc6eaac85aa515/PyYAML-6.0.2-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:68ccc6023a3400877818152ad9a1033e3db8625d899c72eacb5a668902e4d652", size = 751597, upload-time = "2024-08-06T20:32:56.985Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/14/0d/e2c3b43bbce3cf6bd97c840b46088a3031085179e596d4929729d8d68270/PyYAML-6.0.2-cp313-cp313-win32.whl", hash = "sha256:bc2fa7c6b47d6bc618dd7fb02ef6fdedb1090ec036abab80d4681424b84c1183", size = 140527, upload-time = "2024-08-06T20:33:03.001Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fa/de/02b54f42487e3d3c6efb3f89428677074ca7bf43aae402517bc7cca949f3/PyYAML-6.0.2-cp313-cp313-win_amd64.whl", hash = "sha256:8388ee1976c416731879ac16da0aff3f63b286ffdd57cdeb95f3f2e085687563", size = 156446, upload-time = "2024-08-06T20:33:04.33Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/74/d9/323a59d506f12f498c2097488d80d16f4cf965cee1791eab58b56b19f47a/PyYAML-6.0.2-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:24471b829b3bf607e04e88d79542a9d48bb037c2267d7927a874e6c205ca7e9a", size = 183218, upload-time = "2024-08-06T20:33:06.411Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/74/cc/20c34d00f04d785f2028737e2e2a8254e1425102e730fee1d6396f832577/PyYAML-6.0.2-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d7fded462629cfa4b685c5416b949ebad6cec74af5e2d42905d41e257e0869f5", size = 728067, upload-time = "2024-08-06T20:33:07.879Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/20/52/551c69ca1501d21c0de51ddafa8c23a0191ef296ff098e98358f69080577/PyYAML-6.0.2-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:d84a1718ee396f54f3a086ea0a66d8e552b2ab2017ef8b420e92edbc841c352d", size = 757812, upload-time = "2024-08-06T20:33:12.542Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fd/7f/2c3697bba5d4aa5cc2afe81826d73dfae5f049458e44732c7a0938baa673/PyYAML-6.0.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9056c1ecd25795207ad294bcf39f2db3d845767be0ea6e6a34d856f006006083", size = 746531, upload-time = "2024-08-06T20:33:14.391Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8c/ab/6226d3df99900e580091bb44258fde77a8433511a86883bd4681ea19a858/PyYAML-6.0.2-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:82d09873e40955485746739bcb8b4586983670466c23382c19cffecbf1fd8706", size = 800820, upload-time = "2024-08-06T20:33:16.586Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a0/99/a9eb0f3e710c06c5d922026f6736e920d431812ace24aae38228d0d64b04/PyYAML-6.0.2-cp38-cp38-win32.whl", hash = "sha256:43fa96a3ca0d6b1812e01ced1044a003533c47f6ee8aca31724f78e93ccc089a", size = 145514, upload-time = "2024-08-06T20:33:22.414Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/75/8a/ee831ad5fafa4431099aa4e078d4c8efd43cd5e48fbc774641d233b683a9/PyYAML-6.0.2-cp38-cp38-win_amd64.whl", hash = "sha256:01179a4a8559ab5de078078f37e5c1a30d76bb88519906844fd7bdea1b7729ff", size = 162702, upload-time = "2024-08-06T20:33:23.813Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/65/d8/b7a1db13636d7fb7d4ff431593c510c8b8fca920ade06ca8ef20015493c5/PyYAML-6.0.2-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:688ba32a1cffef67fd2e9398a2efebaea461578b0923624778664cc1c914db5d", size = 184777, upload-time = "2024-08-06T20:33:25.896Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0a/02/6ec546cd45143fdf9840b2c6be8d875116a64076218b61d68e12548e5839/PyYAML-6.0.2-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:a8786accb172bd8afb8be14490a16625cbc387036876ab6ba70912730faf8e1f", size = 172318, upload-time = "2024-08-06T20:33:27.212Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0e/9a/8cc68be846c972bda34f6c2a93abb644fb2476f4dcc924d52175786932c9/PyYAML-6.0.2-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d8e03406cac8513435335dbab54c0d385e4a49e4945d2909a581c83647ca0290", size = 720891, upload-time = "2024-08-06T20:33:28.974Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e9/6c/6e1b7f40181bc4805e2e07f4abc10a88ce4648e7e95ff1abe4ae4014a9b2/PyYAML-6.0.2-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:f753120cb8181e736c57ef7636e83f31b9c0d1722c516f7e86cf15b7aa57ff12", size = 722614, upload-time = "2024-08-06T20:33:34.157Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3d/32/e7bd8535d22ea2874cef6a81021ba019474ace0d13a4819c2a4bce79bd6a/PyYAML-6.0.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3b1fdb9dc17f5a7677423d508ab4f243a726dea51fa5e70992e59a7411c89d19", size = 737360, upload-time = "2024-08-06T20:33:35.84Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d7/12/7322c1e30b9be969670b672573d45479edef72c9a0deac3bb2868f5d7469/PyYAML-6.0.2-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:0b69e4ce7a131fe56b7e4d770c67429700908fc0752af059838b1cfb41960e4e", size = 699006, upload-time = "2024-08-06T20:33:37.501Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/82/72/04fcad41ca56491995076630c3ec1e834be241664c0c09a64c9a2589b507/PyYAML-6.0.2-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:a9f8c2e67970f13b16084e04f134610fd1d374bf477b17ec1599185cf611d725", size = 723577, upload-time = "2024-08-06T20:33:39.389Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ed/5e/46168b1f2757f1fcd442bc3029cd8767d88a98c9c05770d8b420948743bb/PyYAML-6.0.2-cp39-cp39-win32.whl", hash = "sha256:6395c297d42274772abc367baaa79683958044e5d3835486c16da75d2a694631", size = 144593, upload-time = "2024-08-06T20:33:46.63Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/19/87/5124b1c1f2412bb95c59ec481eaf936cd32f0fe2a7b16b97b81c4c017a6a/PyYAML-6.0.2-cp39-cp39-win_amd64.whl", hash = "sha256:39693e1f8320ae4f43943590b49779ffb98acb81f788220ea932a6b6c51004d8", size = 162312, upload-time = "2024-08-06T20:33:49.073Z" }, +] + +[[package]] +name = "requests" +version = "2.32.4" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] +dependencies = [ + { name = "certifi", marker = "python_full_version < '3.9'" }, + { name = "charset-normalizer", marker = "python_full_version < '3.9'" }, + { name = "idna", marker = "python_full_version < '3.9'" }, + { name = "urllib3", version = "2.2.3", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e1/0a/929373653770d8a0d7ea76c37de6e41f11eb07559b103b1c02cafb3f7cf8/requests-2.32.4.tar.gz", hash = "sha256:27d0316682c8a29834d3264820024b62a36942083d52caf2f14c0591336d3422", size = 135258, upload-time = "2025-06-09T16:43:07.34Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7c/e4/56027c4a6b4ae70ca9de302488c5ca95ad4a39e190093d6c1a8ace08341b/requests-2.32.4-py3-none-any.whl", hash = "sha256:27babd3cda2a6d50b30443204ee89830707d396671944c998b5975b031ac2b2c", size = 64847, upload-time = "2025-06-09T16:43:05.728Z" }, +] + +[[package]] +name = "requests" +version = "2.32.5" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.13.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'darwin')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'darwin')", + "(python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform == 'darwin')", + "(python_full_version == '3.10.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'darwin')", + "(python_full_version == '3.9.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform == 'darwin')", +] +dependencies = [ + { name = "certifi", marker = "python_full_version >= '3.9'" }, + { name = "charset-normalizer", marker = "python_full_version >= '3.9'" }, + { name = "idna", marker = "python_full_version >= '3.9'" }, + { name = "urllib3", version = "2.5.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c9/74/b3ff8e6c8446842c3f5c837e9c3dfcfe2018ea6ecef224c710c85ef728f4/requests-2.32.5.tar.gz", hash = "sha256:dbba0bac56e100853db0ea71b82b4dfd5fe2bf6d3754a8893c3af500cec7d7cf", size = 134517, upload-time = "2025-08-18T20:46:02.573Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1e/db/4254e3eabe8020b458f1a747140d32277ec7a271daf1d235b70dc0b4e6e3/requests-2.32.5-py3-none-any.whl", hash = "sha256:2462f94637a34fd532264295e186976db0f5d453d1cdd31473c85a6a161affb6", size = 64738, upload-time = "2025-08-18T20:46:00.542Z" }, +] + +[[package]] +name = "rotary-embedding-torch" +version = "0.8.9" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +dependencies = [ + { name = "einops" }, + { name = "torch", version = "2.5.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "torch", version = "2.8.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/71/38/74783585b1f0282fddd3faf1abd6dd20977255c27e31737eced2d7ec05f1/rotary_embedding_torch-0.8.9.tar.gz", hash = "sha256:b213f153cad1d108064d930544fb3af678d56515893d3f869a7a146f87997e3f", size = 7497, upload-time = "2025-07-27T01:26:14.675Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/94/21/05e319ebef4a8194f9c4cc09f9c3213ab846da200fa6e88699cdb7e9679a/rotary_embedding_torch-0.8.9-py3-none-any.whl", hash = "sha256:700e8de9dfbefba5f9117a66652a2520648dcc60136895f0068b6a85347cab02", size = 5872, upload-time = "2025-07-27T01:26:13.496Z" }, +] + +[[package]] +name = "scikit-learn" +version = "1.3.2" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] +dependencies = [ + { name = "joblib", version = "1.4.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "numpy", version = "1.24.4", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "scipy", version = "1.10.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "threadpoolctl", version = "3.5.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/88/00/835e3d280fdd7784e76bdef91dd9487582d7951a7254f59fc8004fc8b213/scikit-learn-1.3.2.tar.gz", hash = "sha256:a2f54c76accc15a34bfb9066e6c7a56c1e7235dda5762b990792330b52ccfb05", size = 7510251, upload-time = "2023-10-23T13:47:55.287Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0d/53/570b55a6e10b8694ac1e3024d2df5cd443f1b4ff6d28430845da8b9019b3/scikit_learn-1.3.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:e326c0eb5cf4d6ba40f93776a20e9a7a69524c4db0757e7ce24ba222471ee8a1", size = 10209999, upload-time = "2023-10-23T13:46:30.373Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/70/d0/50ace22129f79830e3cf682d0a2bd4843ef91573299d43112d52790163a8/scikit_learn-1.3.2-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:535805c2a01ccb40ca4ab7d081d771aea67e535153e35a1fd99418fcedd1648a", size = 9479353, upload-time = "2023-10-23T13:46:34.368Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8f/46/fcc35ed7606c50d3072eae5a107a45cfa5b7f5fa8cc48610edd8cc8e8550/scikit_learn-1.3.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1215e5e58e9880b554b01187b8c9390bf4dc4692eedeaf542d3273f4785e342c", size = 10304705, upload-time = "2023-10-23T13:46:37.868Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d0/0b/26ad95cf0b747be967b15fb71a06f5ac67aba0fd2f9cd174de6edefc4674/scikit_learn-1.3.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0ee107923a623b9f517754ea2f69ea3b62fc898a3641766cb7deb2f2ce450161", size = 10827807, upload-time = "2023-10-23T13:46:41.59Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/69/8a/cf17d6443f5f537e099be81535a56ab68a473f9393fbffda38cd19899fc8/scikit_learn-1.3.2-cp310-cp310-win_amd64.whl", hash = "sha256:35a22e8015048c628ad099da9df5ab3004cdbf81edc75b396fd0cff8699ac58c", size = 9255427, upload-time = "2023-10-23T13:46:44.826Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/08/5d/e5acecd6e99a6b656e42e7a7b18284e2f9c9f512e8ed6979e1e75d25f05f/scikit_learn-1.3.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:6fb6bc98f234fda43163ddbe36df8bcde1d13ee176c6dc9b92bb7d3fc842eb66", size = 10116376, upload-time = "2023-10-23T13:46:48.147Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/40/c6/2e91eefb757822e70d351e02cc38d07c137212ae7c41ac12746415b4860a/scikit_learn-1.3.2-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:18424efee518a1cde7b0b53a422cde2f6625197de6af36da0b57ec502f126157", size = 9383415, upload-time = "2023-10-23T13:46:51.324Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fa/fd/b3637639e73bb72b12803c5245f2a7299e09b2acd85a0f23937c53369a1c/scikit_learn-1.3.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3271552a5eb16f208a6f7f617b8cc6d1f137b52c8a1ef8edf547db0259b2c9fb", size = 10279163, upload-time = "2023-10-23T13:46:54.642Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0c/2a/d3ff6091406bc2207e0adb832ebd15e40ac685811c7e2e3b432bfd969b71/scikit_learn-1.3.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fc4144a5004a676d5022b798d9e573b05139e77f271253a4703eed295bde0433", size = 10884422, upload-time = "2023-10-23T13:46:58.087Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4e/ba/ce9bd1cd4953336a0e213b29cb80bb11816f2a93de8c99f88ef0b446ad0c/scikit_learn-1.3.2-cp311-cp311-win_amd64.whl", hash = "sha256:67f37d708f042a9b8d59551cf94d30431e01374e00dc2645fa186059c6c5d78b", size = 9207060, upload-time = "2023-10-23T13:47:00.948Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/26/7e/2c3b82c8c29aa384c8bf859740419278627d2cdd0050db503c8840e72477/scikit_learn-1.3.2-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:8db94cd8a2e038b37a80a04df8783e09caac77cbe052146432e67800e430c028", size = 9979322, upload-time = "2023-10-23T13:47:03.977Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/cf/fc/6c52ffeb587259b6b893b7cac268f1eb1b5426bcce1aa20e53523bfe6944/scikit_learn-1.3.2-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:61a6efd384258789aa89415a410dcdb39a50e19d3d8410bd29be365bcdd512d5", size = 9270688, upload-time = "2023-10-23T13:47:07.316Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e5/a7/6f4ae76f72ae9de162b97acbf1f53acbe404c555f968d13da21e4112a002/scikit_learn-1.3.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:cb06f8dce3f5ddc5dee1715a9b9f19f20d295bed8e3cd4fa51e1d050347de525", size = 10280398, upload-time = "2023-10-23T13:47:10.796Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5d/b7/ee35904c07a0666784349529412fbb9814a56382b650d30fd9d6be5e5054/scikit_learn-1.3.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5b2de18d86f630d68fe1f87af690d451388bb186480afc719e5f770590c2ef6c", size = 10796478, upload-time = "2023-10-23T13:47:14.077Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fe/6b/db949ed5ac367987b1f250f070f340b7715d22f0c9c965bdf07de6ca75a3/scikit_learn-1.3.2-cp312-cp312-win_amd64.whl", hash = "sha256:0402638c9a7c219ee52c94cbebc8fcb5eb9fe9c773717965c1f4185588ad3107", size = 9133979, upload-time = "2023-10-23T13:47:17.389Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e3/52/fd60b0b022af41fbf3463587ddc719288f0f2d4e46603ab3184996cd5f04/scikit_learn-1.3.2-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:a19f90f95ba93c1a7f7924906d0576a84da7f3b2282ac3bfb7a08a32801add93", size = 10064879, upload-time = "2023-10-23T13:47:21.392Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a4/62/92e9cec3deca8b45abf62dd8f6469d688b3f28b9c170809fcc46f110b523/scikit_learn-1.3.2-cp38-cp38-macosx_12_0_arm64.whl", hash = "sha256:b8692e395a03a60cd927125eef3a8e3424d86dde9b2370d544f0ea35f78a8073", size = 9373934, upload-time = "2023-10-23T13:47:24.645Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/49/81/91585dc83ec81dcd52e934f6708bf350b06949d8bfa13bf3b711b851c3f4/scikit_learn-1.3.2-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:15e1e94cc23d04d39da797ee34236ce2375ddea158b10bee3c343647d615581d", size = 10499159, upload-time = "2023-10-23T13:47:28.41Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3f/48/6fdd99f5717045f9984616b5c2ec683d6286d30c0ac234563062132b83ab/scikit_learn-1.3.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:785a2213086b7b1abf037aeadbbd6d67159feb3e30263434139c98425e3dcfcf", size = 11067392, upload-time = "2023-10-23T13:47:32.087Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/52/2d/ad6928a578c78bb0e44e34a5a922818b14c56716b81d145924f1f291416f/scikit_learn-1.3.2-cp38-cp38-win_amd64.whl", hash = "sha256:64381066f8aa63c2710e6b56edc9f0894cc7bf59bd71b8ce5613a4559b6145e0", size = 9257871, upload-time = "2023-10-23T13:47:36.142Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f8/67/584acfc492ae1bd293d80c7a8c57ba7456e4e415c64869b7c240679eaf78/scikit_learn-1.3.2-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:6c43290337f7a4b969d207e620658372ba3c1ffb611f8bc2b6f031dc5c6d1d03", size = 10232286, upload-time = "2023-10-23T13:47:39.434Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/20/0f/51e3ccdc87c25e2e33bf7962249ff8c5ab1d6aed0144fb003348ce8bd352/scikit_learn-1.3.2-cp39-cp39-macosx_12_0_arm64.whl", hash = "sha256:dc9002fc200bed597d5d34e90c752b74df516d592db162f756cc52836b38fe0e", size = 9504918, upload-time = "2023-10-23T13:47:42.679Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/61/2e/5bbf3c9689d2911b65297fb5861c4257e54c797b3158c9fca8a5c576644b/scikit_learn-1.3.2-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1d08ada33e955c54355d909b9c06a4789a729977f165b8bae6f225ff0a60ec4a", size = 10358127, upload-time = "2023-10-23T13:47:45.96Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/25/89/dce01a35d354159dcc901e3c7e7eb3fe98de5cb3639c6cd39518d8830caa/scikit_learn-1.3.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:763f0ae4b79b0ff9cca0bf3716bcc9915bdacff3cebea15ec79652d1cc4fa5c9", size = 10890482, upload-time = "2023-10-23T13:47:49.046Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1c/49/30ffcac5af06d08dfdd27da322ce31a373b733711bb272941877c1e4794a/scikit_learn-1.3.2-cp39-cp39-win_amd64.whl", hash = "sha256:ed932ea780517b00dae7431e031faae6b49b20eb6950918eb83bd043237950e0", size = 9331050, upload-time = "2023-10-23T13:47:52.246Z" }, +] + +[[package]] +name = "scikit-learn" +version = "1.6.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform == 'darwin')", +] +dependencies = [ + { name = "joblib", version = "1.5.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, + { name = "numpy", version = "2.0.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, + { name = "scipy", version = "1.13.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, + { name = "threadpoolctl", version = "3.6.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/9e/a5/4ae3b3a0755f7b35a280ac90b28817d1f380318973cff14075ab41ef50d9/scikit_learn-1.6.1.tar.gz", hash = "sha256:b4fc2525eca2c69a59260f583c56a7557c6ccdf8deafdba6e060f94c1c59738e", size = 7068312, upload-time = "2025-01-10T08:07:55.348Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2e/3a/f4597eb41049110b21ebcbb0bcb43e4035017545daa5eedcfeb45c08b9c5/scikit_learn-1.6.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:d056391530ccd1e501056160e3c9673b4da4805eb67eb2bdf4e983e1f9c9204e", size = 12067702, upload-time = "2025-01-10T08:05:56.515Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/37/19/0423e5e1fd1c6ec5be2352ba05a537a473c1677f8188b9306097d684b327/scikit_learn-1.6.1-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:0c8d036eb937dbb568c6242fa598d551d88fb4399c0344d95c001980ec1c7d36", size = 11112765, upload-time = "2025-01-10T08:06:00.272Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/70/95/d5cb2297a835b0f5fc9a77042b0a2d029866379091ab8b3f52cc62277808/scikit_learn-1.6.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8634c4bd21a2a813e0a7e3900464e6d593162a29dd35d25bdf0103b3fce60ed5", size = 12643991, upload-time = "2025-01-10T08:06:04.813Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b7/91/ab3c697188f224d658969f678be86b0968ccc52774c8ab4a86a07be13c25/scikit_learn-1.6.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:775da975a471c4f6f467725dff0ced5c7ac7bda5e9316b260225b48475279a1b", size = 13497182, upload-time = "2025-01-10T08:06:08.42Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/17/04/d5d556b6c88886c092cc989433b2bab62488e0f0dafe616a1d5c9cb0efb1/scikit_learn-1.6.1-cp310-cp310-win_amd64.whl", hash = "sha256:8a600c31592bd7dab31e1c61b9bbd6dea1b3433e67d264d17ce1017dbdce8002", size = 11125517, upload-time = "2025-01-10T08:06:12.783Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6c/2a/e291c29670795406a824567d1dfc91db7b699799a002fdaa452bceea8f6e/scikit_learn-1.6.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:72abc587c75234935e97d09aa4913a82f7b03ee0b74111dcc2881cba3c5a7b33", size = 12102620, upload-time = "2025-01-10T08:06:16.675Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/25/92/ee1d7a00bb6b8c55755d4984fd82608603a3cc59959245068ce32e7fb808/scikit_learn-1.6.1-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:b3b00cdc8f1317b5f33191df1386c0befd16625f49d979fe77a8d44cae82410d", size = 11116234, upload-time = "2025-01-10T08:06:21.83Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/30/cd/ed4399485ef364bb25f388ab438e3724e60dc218c547a407b6e90ccccaef/scikit_learn-1.6.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:dc4765af3386811c3ca21638f63b9cf5ecf66261cc4815c1db3f1e7dc7b79db2", size = 12592155, upload-time = "2025-01-10T08:06:27.309Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a8/f3/62fc9a5a659bb58a03cdd7e258956a5824bdc9b4bb3c5d932f55880be569/scikit_learn-1.6.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:25fc636bdaf1cc2f4a124a116312d837148b5e10872147bdaf4887926b8c03d8", size = 13497069, upload-time = "2025-01-10T08:06:32.515Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a1/a6/c5b78606743a1f28eae8f11973de6613a5ee87366796583fb74c67d54939/scikit_learn-1.6.1-cp311-cp311-win_amd64.whl", hash = "sha256:fa909b1a36e000a03c382aade0bd2063fd5680ff8b8e501660c0f59f021a6415", size = 11139809, upload-time = "2025-01-10T08:06:35.514Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0a/18/c797c9b8c10380d05616db3bfb48e2a3358c767affd0857d56c2eb501caa/scikit_learn-1.6.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:926f207c804104677af4857b2c609940b743d04c4c35ce0ddc8ff4f053cddc1b", size = 12104516, upload-time = "2025-01-10T08:06:40.009Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c4/b7/2e35f8e289ab70108f8cbb2e7a2208f0575dc704749721286519dcf35f6f/scikit_learn-1.6.1-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:2c2cae262064e6a9b77eee1c8e768fc46aa0b8338c6a8297b9b6759720ec0ff2", size = 11167837, upload-time = "2025-01-10T08:06:43.305Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a4/f6/ff7beaeb644bcad72bcfd5a03ff36d32ee4e53a8b29a639f11bcb65d06cd/scikit_learn-1.6.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1061b7c028a8663fb9a1a1baf9317b64a257fcb036dae5c8752b2abef31d136f", size = 12253728, upload-time = "2025-01-10T08:06:47.618Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/29/7a/8bce8968883e9465de20be15542f4c7e221952441727c4dad24d534c6d99/scikit_learn-1.6.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2e69fab4ebfc9c9b580a7a80111b43d214ab06250f8a7ef590a4edf72464dd86", size = 13147700, upload-time = "2025-01-10T08:06:50.888Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/62/27/585859e72e117fe861c2079bcba35591a84f801e21bc1ab85bce6ce60305/scikit_learn-1.6.1-cp312-cp312-win_amd64.whl", hash = "sha256:70b1d7e85b1c96383f872a519b3375f92f14731e279a7b4c6cfd650cf5dffc52", size = 11110613, upload-time = "2025-01-10T08:06:54.115Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2e/59/8eb1872ca87009bdcdb7f3cdc679ad557b992c12f4b61f9250659e592c63/scikit_learn-1.6.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:2ffa1e9e25b3d93990e74a4be2c2fc61ee5af85811562f1288d5d055880c4322", size = 12010001, upload-time = "2025-01-10T08:06:58.613Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/9d/05/f2fc4effc5b32e525408524c982c468c29d22f828834f0625c5ef3d601be/scikit_learn-1.6.1-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:dc5cf3d68c5a20ad6d571584c0750ec641cc46aeef1c1507be51300e6003a7e1", size = 11096360, upload-time = "2025-01-10T08:07:01.556Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c8/e4/4195d52cf4f113573fb8ebc44ed5a81bd511a92c0228889125fac2f4c3d1/scikit_learn-1.6.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c06beb2e839ecc641366000ca84f3cf6fa9faa1777e29cf0c04be6e4d096a348", size = 12209004, upload-time = "2025-01-10T08:07:06.931Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/94/be/47e16cdd1e7fcf97d95b3cb08bde1abb13e627861af427a3651fcb80b517/scikit_learn-1.6.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e8ca8cb270fee8f1f76fa9bfd5c3507d60c6438bbee5687f81042e2bb98e5a97", size = 13171776, upload-time = "2025-01-10T08:07:11.715Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/34/b0/ca92b90859070a1487827dbc672f998da95ce83edce1270fc23f96f1f61a/scikit_learn-1.6.1-cp313-cp313-win_amd64.whl", hash = "sha256:7a1c43c8ec9fde528d664d947dc4c0789be4077a3647f232869f41d9bf50e0fb", size = 11071865, upload-time = "2025-01-10T08:07:16.088Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/12/ae/993b0fb24a356e71e9a894e42b8a9eec528d4c70217353a1cd7a48bc25d4/scikit_learn-1.6.1-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:a17c1dea1d56dcda2fac315712f3651a1fea86565b64b48fa1bc090249cbf236", size = 11955804, upload-time = "2025-01-10T08:07:20.385Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d6/54/32fa2ee591af44507eac86406fa6bba968d1eb22831494470d0a2e4a1eb1/scikit_learn-1.6.1-cp313-cp313t-macosx_12_0_arm64.whl", hash = "sha256:6a7aa5f9908f0f28f4edaa6963c0a6183f1911e63a69aa03782f0d924c830a35", size = 11100530, upload-time = "2025-01-10T08:07:23.675Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3f/58/55856da1adec655bdce77b502e94a267bf40a8c0b89f8622837f89503b5a/scikit_learn-1.6.1-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0650e730afb87402baa88afbf31c07b84c98272622aaba002559b614600ca691", size = 12433852, upload-time = "2025-01-10T08:07:26.817Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ff/4f/c83853af13901a574f8f13b645467285a48940f185b690936bb700a50863/scikit_learn-1.6.1-cp313-cp313t-win_amd64.whl", hash = "sha256:3f59fe08dc03ea158605170eb52b22a105f238a5d512c4470ddeca71feae8e5f", size = 11337256, upload-time = "2025-01-10T08:07:31.084Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d2/37/b305b759cc65829fe1b8853ff3e308b12cdd9d8884aa27840835560f2b42/scikit_learn-1.6.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:6849dd3234e87f55dce1db34c89a810b489ead832aaf4d4550b7ea85628be6c1", size = 12101868, upload-time = "2025-01-10T08:07:34.189Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/83/74/f64379a4ed5879d9db744fe37cfe1978c07c66684d2439c3060d19a536d8/scikit_learn-1.6.1-cp39-cp39-macosx_12_0_arm64.whl", hash = "sha256:e7be3fa5d2eb9be7d77c3734ff1d599151bb523674be9b834e8da6abe132f44e", size = 11144062, upload-time = "2025-01-10T08:07:37.67Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fd/dc/d5457e03dc9c971ce2b0d750e33148dd060fefb8b7dc71acd6054e4bb51b/scikit_learn-1.6.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:44a17798172df1d3c1065e8fcf9019183f06c87609b49a124ebdf57ae6cb0107", size = 12693173, upload-time = "2025-01-10T08:07:42.713Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/79/35/b1d2188967c3204c78fa79c9263668cf1b98060e8e58d1a730fe5b2317bb/scikit_learn-1.6.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b8b7a3b86e411e4bce21186e1c180d792f3d99223dcfa3b4f597ecc92fa1a422", size = 13518605, upload-time = "2025-01-10T08:07:46.551Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fb/d8/8d603bdd26601f4b07e2363032b8565ab82eb857f93d86d0f7956fcf4523/scikit_learn-1.6.1-cp39-cp39-win_amd64.whl", hash = "sha256:7a73d457070e3318e32bdb3aa79a8d990474f19035464dfd8bede2883ab5dc3b", size = 11155078, upload-time = "2025-01-10T08:07:51.376Z" }, +] + +[[package]] +name = "scikit-learn" +version = "1.7.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.13.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'darwin')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'darwin')", + "(python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform == 'darwin')", + "(python_full_version == '3.10.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'darwin')", +] +dependencies = [ + { name = "joblib", version = "1.5.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.10'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.10'" }, + { name = "scipy", version = "1.15.3", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.10.*'" }, + { name = "scipy", version = "1.16.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "threadpoolctl", version = "3.6.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.10'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/41/84/5f4af978fff619706b8961accac84780a6d298d82a8873446f72edb4ead0/scikit_learn-1.7.1.tar.gz", hash = "sha256:24b3f1e976a4665aa74ee0fcaac2b8fccc6ae77c8e07ab25da3ba6d3292b9802", size = 7190445, upload-time = "2025-07-18T08:01:54.5Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/74/88/0dd5be14ef19f2d80a77780be35a33aa94e8a3b3223d80bee8892a7832b4/scikit_learn-1.7.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:406204dd4004f0517f0b23cf4b28c6245cbd51ab1b6b78153bc784def214946d", size = 9338868, upload-time = "2025-07-18T08:01:00.25Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fd/52/3056b6adb1ac58a0bc335fc2ed2fcf599974d908855e8cb0ca55f797593c/scikit_learn-1.7.1-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:16af2e44164f05d04337fd1fc3ae7c4ea61fd9b0d527e22665346336920fe0e1", size = 8655943, upload-time = "2025-07-18T08:01:02.974Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fb/a4/e488acdece6d413f370a9589a7193dac79cd486b2e418d3276d6ea0b9305/scikit_learn-1.7.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2f2e78e56a40c7587dea9a28dc4a49500fa2ead366869418c66f0fd75b80885c", size = 9652056, upload-time = "2025-07-18T08:01:04.978Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/18/41/bceacec1285b94eb9e4659b24db46c23346d7e22cf258d63419eb5dec6f7/scikit_learn-1.7.1-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b62b76ad408a821475b43b7bb90a9b1c9a4d8d125d505c2df0539f06d6e631b1", size = 9473691, upload-time = "2025-07-18T08:01:07.006Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/12/7b/e1ae4b7e1dd85c4ca2694ff9cc4a9690970fd6150d81b975e6c5c6f8ee7c/scikit_learn-1.7.1-cp310-cp310-win_amd64.whl", hash = "sha256:9963b065677a4ce295e8ccdee80a1dd62b37249e667095039adcd5bce6e90deb", size = 8900873, upload-time = "2025-07-18T08:01:09.332Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b4/bd/a23177930abd81b96daffa30ef9c54ddbf544d3226b8788ce4c3ef1067b4/scikit_learn-1.7.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:90c8494ea23e24c0fb371afc474618c1019dc152ce4a10e4607e62196113851b", size = 9334838, upload-time = "2025-07-18T08:01:11.239Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8d/a1/d3a7628630a711e2ac0d1a482910da174b629f44e7dd8cfcd6924a4ef81a/scikit_learn-1.7.1-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:bb870c0daf3bf3be145ec51df8ac84720d9972170786601039f024bf6d61a518", size = 8651241, upload-time = "2025-07-18T08:01:13.234Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/26/92/85ec172418f39474c1cd0221d611345d4f433fc4ee2fc68e01f524ccc4e4/scikit_learn-1.7.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:40daccd1b5623f39e8943ab39735cadf0bdce80e67cdca2adcb5426e987320a8", size = 9718677, upload-time = "2025-07-18T08:01:15.649Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/df/ce/abdb1dcbb1d2b66168ec43b23ee0cee356b4cc4100ddee3943934ebf1480/scikit_learn-1.7.1-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:30d1f413cfc0aa5a99132a554f1d80517563c34a9d3e7c118fde2d273c6fe0f7", size = 9511189, upload-time = "2025-07-18T08:01:18.013Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b2/3b/47b5eaee01ef2b5a80ba3f7f6ecf79587cb458690857d4777bfd77371c6f/scikit_learn-1.7.1-cp311-cp311-win_amd64.whl", hash = "sha256:c711d652829a1805a95d7fe96654604a8f16eab5a9e9ad87b3e60173415cb650", size = 8914794, upload-time = "2025-07-18T08:01:20.357Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/cb/16/57f176585b35ed865f51b04117947fe20f130f78940c6477b6d66279c9c2/scikit_learn-1.7.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:3cee419b49b5bbae8796ecd690f97aa412ef1674410c23fc3257c6b8b85b8087", size = 9260431, upload-time = "2025-07-18T08:01:22.77Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/67/4e/899317092f5efcab0e9bc929e3391341cec8fb0e816c4789686770024580/scikit_learn-1.7.1-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:2fd8b8d35817b0d9ebf0b576f7d5ffbbabdb55536b0655a8aaae629d7ffd2e1f", size = 8637191, upload-time = "2025-07-18T08:01:24.731Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f3/1b/998312db6d361ded1dd56b457ada371a8d8d77ca2195a7d18fd8a1736f21/scikit_learn-1.7.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:588410fa19a96a69763202f1d6b7b91d5d7a5d73be36e189bc6396bfb355bd87", size = 9486346, upload-time = "2025-07-18T08:01:26.713Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ad/09/a2aa0b4e644e5c4ede7006748f24e72863ba2ae71897fecfd832afea01b4/scikit_learn-1.7.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e3142f0abe1ad1d1c31a2ae987621e41f6b578144a911ff4ac94781a583adad7", size = 9290988, upload-time = "2025-07-18T08:01:28.938Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/15/fa/c61a787e35f05f17fc10523f567677ec4eeee5f95aa4798dbbbcd9625617/scikit_learn-1.7.1-cp312-cp312-win_amd64.whl", hash = "sha256:3ddd9092c1bd469acab337d87930067c87eac6bd544f8d5027430983f1e1ae88", size = 8735568, upload-time = "2025-07-18T08:01:30.936Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/52/f8/e0533303f318a0f37b88300d21f79b6ac067188d4824f1047a37214ab718/scikit_learn-1.7.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:b7839687fa46d02e01035ad775982f2470be2668e13ddd151f0f55a5bf123bae", size = 9213143, upload-time = "2025-07-18T08:01:32.942Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/71/f3/f1df377d1bdfc3e3e2adc9c119c238b182293e6740df4cbeac6de2cc3e23/scikit_learn-1.7.1-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:a10f276639195a96c86aa572ee0698ad64ee939a7b042060b98bd1930c261d10", size = 8591977, upload-time = "2025-07-18T08:01:34.967Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/99/72/c86a4cd867816350fe8dee13f30222340b9cd6b96173955819a5561810c5/scikit_learn-1.7.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:13679981fdaebc10cc4c13c43344416a86fcbc61449cb3e6517e1df9d12c8309", size = 9436142, upload-time = "2025-07-18T08:01:37.397Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e8/66/277967b29bd297538dc7a6ecfb1a7dce751beabd0d7f7a2233be7a4f7832/scikit_learn-1.7.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4f1262883c6a63f067a980a8cdd2d2e7f2513dddcef6a9eaada6416a7a7cbe43", size = 9282996, upload-time = "2025-07-18T08:01:39.721Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e2/47/9291cfa1db1dae9880420d1e07dbc7e8dd4a7cdbc42eaba22512e6bde958/scikit_learn-1.7.1-cp313-cp313-win_amd64.whl", hash = "sha256:ca6d31fb10e04d50bfd2b50d66744729dbb512d4efd0223b864e2fdbfc4cee11", size = 8707418, upload-time = "2025-07-18T08:01:42.124Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/61/95/45726819beccdaa34d3362ea9b2ff9f2b5d3b8bf721bd632675870308ceb/scikit_learn-1.7.1-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:781674d096303cfe3d351ae6963ff7c958db61cde3421cd490e3a5a58f2a94ae", size = 9561466, upload-time = "2025-07-18T08:01:44.195Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ee/1c/6f4b3344805de783d20a51eb24d4c9ad4b11a7f75c1801e6ec6d777361fd/scikit_learn-1.7.1-cp313-cp313t-macosx_12_0_arm64.whl", hash = "sha256:10679f7f125fe7ecd5fad37dd1aa2daae7e3ad8df7f3eefa08901b8254b3e12c", size = 9040467, upload-time = "2025-07-18T08:01:46.671Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6f/80/abe18fe471af9f1d181904203d62697998b27d9b62124cd281d740ded2f9/scikit_learn-1.7.1-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1f812729e38c8cb37f760dce71a9b83ccfb04f59b3dca7c6079dcdc60544fa9e", size = 9532052, upload-time = "2025-07-18T08:01:48.676Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/14/82/b21aa1e0c4cee7e74864d3a5a721ab8fcae5ca55033cb6263dca297ed35b/scikit_learn-1.7.1-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:88e1a20131cf741b84b89567e1717f27a2ced228e0f29103426102bc2e3b8ef7", size = 9361575, upload-time = "2025-07-18T08:01:50.639Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f2/20/f4777fcd5627dc6695fa6b92179d0edb7a3ac1b91bcd9a1c7f64fa7ade23/scikit_learn-1.7.1-cp313-cp313t-win_amd64.whl", hash = "sha256:b1bd1d919210b6a10b7554b717c9000b5485aa95a1d0f177ae0d7ee8ec750da5", size = 9277310, upload-time = "2025-07-18T08:01:52.547Z" }, +] + +[[package]] +name = "scipy" +version = "1.10.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] +dependencies = [ + { name = "numpy", version = "1.24.4", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/84/a9/2bf119f3f9cff1f376f924e39cfae18dec92a1514784046d185731301281/scipy-1.10.1.tar.gz", hash = "sha256:2cf9dfb80a7b4589ba4c40ce7588986d6d5cebc5457cad2c2880f6bc2d42f3a5", size = 42407997, upload-time = "2023-02-19T21:20:13.395Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0a/ac/b1f1bbf7b01d96495f35be003b881f10f85bf6559efb6e9578da832c2140/scipy-1.10.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:e7354fd7527a4b0377ce55f286805b34e8c54b91be865bac273f527e1b839019", size = 35093243, upload-time = "2023-02-19T20:33:55.754Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ea/e5/452086ebed676ce4000ceb5eeeb0ee4f8c6f67c7e70fb9323a370ff95c1f/scipy-1.10.1-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:4b3f429188c66603a1a5c549fb414e4d3bdc2a24792e061ffbd607d3d75fd84e", size = 28772969, upload-time = "2023-02-19T20:34:39.318Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/04/0b/a1b119c869b79a2ab459b7f9fd7e2dea75a9c7d432e64e915e75586bd00b/scipy-1.10.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1553b5dcddd64ba9a0d95355e63fe6c3fc303a8fd77c7bc91e77d61363f7433f", size = 30886961, upload-time = "2023-02-19T20:35:33.724Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1f/4b/3bacad9a166350cb2e518cea80ab891016933cc1653f15c90279512c5fa9/scipy-1.10.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4c0ff64b06b10e35215abce517252b375e580a6125fd5fdf6421b98efbefb2d2", size = 34422544, upload-time = "2023-02-19T20:37:03.859Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ec/e3/b06ac3738bf365e89710205a471abe7dceec672a51c244b469bc5d1291c7/scipy-1.10.1-cp310-cp310-win_amd64.whl", hash = "sha256:fae8a7b898c42dffe3f7361c40d5952b6bf32d10c4569098d276b4c547905ee1", size = 42484848, upload-time = "2023-02-19T20:39:09.467Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e7/53/053cd3669be0d474deae8fe5f757bff4c4f480b8a410231e0631c068873d/scipy-1.10.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:0f1564ea217e82c1bbe75ddf7285ba0709ecd503f048cb1236ae9995f64217bd", size = 35003170, upload-time = "2023-02-19T20:40:53.274Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0d/3e/d05b9de83677195886fb79844fcca19609a538db63b1790fa373155bc3cf/scipy-1.10.1-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:d925fa1c81b772882aa55bcc10bf88324dadb66ff85d548c71515f6689c6dac5", size = 28717513, upload-time = "2023-02-19T20:42:20.82Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a5/3d/b69746c50e44893da57a68457da3d7e5bb75f6a37fbace3769b70d017488/scipy-1.10.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:aaea0a6be54462ec027de54fca511540980d1e9eea68b2d5c1dbfe084797be35", size = 30687257, upload-time = "2023-02-19T20:43:48.139Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/21/cd/fe2d4af234b80dc08c911ce63fdaee5badcdde3e9bcd9a68884580652ef0/scipy-1.10.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:15a35c4242ec5f292c3dd364a7c71a61be87a3d4ddcc693372813c0b73c9af1d", size = 34124096, upload-time = "2023-02-19T20:45:27.415Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/65/76/903324159e4a3566e518c558aeb21571d642f781d842d8dd0fd9c6b0645a/scipy-1.10.1-cp311-cp311-win_amd64.whl", hash = "sha256:43b8e0bcb877faf0abfb613d51026cd5cc78918e9530e375727bf0625c82788f", size = 42238704, upload-time = "2023-02-19T20:47:26.366Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a0/e3/37508a11dae501349d7c16e4dd18c706a023629eedc650ee094593887a89/scipy-1.10.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:5678f88c68ea866ed9ebe3a989091088553ba12c6090244fdae3e467b1139c35", size = 35041063, upload-time = "2023-02-19T20:49:02.296Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/93/4a/50c436de1353cce8b66b26e49a687f10b91fe7465bf34e4565d810153003/scipy-1.10.1-cp38-cp38-macosx_12_0_arm64.whl", hash = "sha256:39becb03541f9e58243f4197584286e339029e8908c46f7221abeea4b749fa88", size = 28797694, upload-time = "2023-02-19T20:50:19.381Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d2/b5/ff61b79ad0ebd15d87ade10e0f4e80114dd89fac34a5efade39e99048c91/scipy-1.10.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bce5869c8d68cf383ce240e44c1d9ae7c06078a9396df68ce88a1230f93a30c1", size = 31024657, upload-time = "2023-02-19T20:51:49.175Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/69/f0/fb07a9548e48b687b8bf2fa81d71aba9cfc548d365046ca1c791e24db99d/scipy-1.10.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:07c3457ce0b3ad5124f98a86533106b643dd811dd61b548e78cf4c8786652f6f", size = 34540352, upload-time = "2023-02-19T20:53:30.821Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/32/8e/7f403535ddf826348c9b8417791e28712019962f7e90ff845896d6325d09/scipy-1.10.1-cp38-cp38-win_amd64.whl", hash = "sha256:049a8bbf0ad95277ffba9b3b7d23e5369cc39e66406d60422c8cfef40ccc8415", size = 42215036, upload-time = "2023-02-19T20:55:09.639Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d9/7d/78b8035bc93c869b9f17261c87aae97a9cdb937f65f0d453c2831aa172fc/scipy-1.10.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:cd9f1027ff30d90618914a64ca9b1a77a431159df0e2a195d8a9e8a04c78abf9", size = 35158611, upload-time = "2023-02-19T20:56:02.715Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e7/f0/55d81813b1a4cb79ce7dc8290eac083bf38bfb36e1ada94ea13b7b1a5f79/scipy-1.10.1-cp39-cp39-macosx_12_0_arm64.whl", hash = "sha256:79c8e5a6c6ffaf3a2262ef1be1e108a035cf4f05c14df56057b64acc5bebffb6", size = 28902591, upload-time = "2023-02-19T20:56:45.728Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/77/d1/722c457b319eed1d642e0a14c9be37eb475f0e6ed1f3401fa480d5d6d36e/scipy-1.10.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:51af417a000d2dbe1ec6c372dfe688e041a7084da4fdd350aeb139bd3fb55353", size = 30960654, upload-time = "2023-02-19T20:57:32.091Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5d/30/b2a2a5bf1a3beefb7609fb871dcc6aef7217c69cef19a4631b7ab5622a8a/scipy-1.10.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1b4735d6c28aad3cdcf52117e0e91d6b39acd4272f3f5cd9907c24ee931ad601", size = 34458863, upload-time = "2023-02-19T20:58:23.601Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/35/20/0ec6246bbb43d18650c9a7cad6602e1a84fd8f9564a9b84cc5faf1e037d0/scipy-1.10.1-cp39-cp39-win_amd64.whl", hash = "sha256:7ff7f37b1bf4417baca958d254e8e2875d0cc23aaadbe65b3d5b3077b0eb23ea", size = 42509516, upload-time = "2023-02-19T20:59:26.296Z" }, +] + +[[package]] +name = "scipy" +version = "1.13.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform == 'darwin')", +] +dependencies = [ + { name = "numpy", version = "2.0.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ae/00/48c2f661e2816ccf2ecd77982f6605b2950afe60f60a52b4cbbc2504aa8f/scipy-1.13.1.tar.gz", hash = "sha256:095a87a0312b08dfd6a6155cbbd310a8c51800fc931b8c0b84003014b874ed3c", size = 57210720, upload-time = "2024-05-23T03:29:26.079Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/33/59/41b2529908c002ade869623b87eecff3e11e3ce62e996d0bdcb536984187/scipy-1.13.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:20335853b85e9a49ff7572ab453794298bcf0354d8068c5f6775a0eabf350aca", size = 39328076, upload-time = "2024-05-23T03:19:01.687Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d5/33/f1307601f492f764062ce7dd471a14750f3360e33cd0f8c614dae208492c/scipy-1.13.1-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:d605e9c23906d1994f55ace80e0125c587f96c020037ea6aa98d01b4bd2e222f", size = 30306232, upload-time = "2024-05-23T03:19:09.089Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c0/66/9cd4f501dd5ea03e4a4572ecd874936d0da296bd04d1c45ae1a4a75d9c3a/scipy-1.13.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:cfa31f1def5c819b19ecc3a8b52d28ffdcc7ed52bb20c9a7589669dd3c250989", size = 33743202, upload-time = "2024-05-23T03:19:15.138Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a3/ba/7255e5dc82a65adbe83771c72f384d99c43063648456796436c9a5585ec3/scipy-1.13.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f26264b282b9da0952a024ae34710c2aff7d27480ee91a2e82b7b7073c24722f", size = 38577335, upload-time = "2024-05-23T03:19:21.984Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/49/a5/bb9ded8326e9f0cdfdc412eeda1054b914dfea952bda2097d174f8832cc0/scipy-1.13.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:eccfa1906eacc02de42d70ef4aecea45415f5be17e72b61bafcfd329bdc52e94", size = 38820728, upload-time = "2024-05-23T03:19:28.225Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/12/30/df7a8fcc08f9b4a83f5f27cfaaa7d43f9a2d2ad0b6562cced433e5b04e31/scipy-1.13.1-cp310-cp310-win_amd64.whl", hash = "sha256:2831f0dc9c5ea9edd6e51e6e769b655f08ec6db6e2e10f86ef39bd32eb11da54", size = 46210588, upload-time = "2024-05-23T03:19:35.661Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b4/15/4a4bb1b15bbd2cd2786c4f46e76b871b28799b67891f23f455323a0cdcfb/scipy-1.13.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:27e52b09c0d3a1d5b63e1105f24177e544a222b43611aaf5bc44d4a0979e32f9", size = 39333805, upload-time = "2024-05-23T03:19:43.081Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ba/92/42476de1af309c27710004f5cdebc27bec62c204db42e05b23a302cb0c9a/scipy-1.13.1-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:54f430b00f0133e2224c3ba42b805bfd0086fe488835effa33fa291561932326", size = 30317687, upload-time = "2024-05-23T03:19:48.799Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/80/ba/8be64fe225360a4beb6840f3cbee494c107c0887f33350d0a47d55400b01/scipy-1.13.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e89369d27f9e7b0884ae559a3a956e77c02114cc60a6058b4e5011572eea9299", size = 33694638, upload-time = "2024-05-23T03:19:55.104Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/36/07/035d22ff9795129c5a847c64cb43c1fa9188826b59344fee28a3ab02e283/scipy-1.13.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a78b4b3345f1b6f68a763c6e25c0c9a23a9fd0f39f5f3d200efe8feda560a5fa", size = 38569931, upload-time = "2024-05-23T03:20:01.82Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d9/10/f9b43de37e5ed91facc0cfff31d45ed0104f359e4f9a68416cbf4e790241/scipy-1.13.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:45484bee6d65633752c490404513b9ef02475b4284c4cfab0ef946def50b3f59", size = 38838145, upload-time = "2024-05-23T03:20:09.173Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4a/48/4513a1a5623a23e95f94abd675ed91cfb19989c58e9f6f7d03990f6caf3d/scipy-1.13.1-cp311-cp311-win_amd64.whl", hash = "sha256:5713f62f781eebd8d597eb3f88b8bf9274e79eeabf63afb4a737abc6c84ad37b", size = 46196227, upload-time = "2024-05-23T03:20:16.433Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f2/7b/fb6b46fbee30fc7051913068758414f2721003a89dd9a707ad49174e3843/scipy-1.13.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:5d72782f39716b2b3509cd7c33cdc08c96f2f4d2b06d51e52fb45a19ca0c86a1", size = 39357301, upload-time = "2024-05-23T03:20:23.538Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/dc/5a/2043a3bde1443d94014aaa41e0b50c39d046dda8360abd3b2a1d3f79907d/scipy-1.13.1-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:017367484ce5498445aade74b1d5ab377acdc65e27095155e448c88497755a5d", size = 30363348, upload-time = "2024-05-23T03:20:29.885Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e7/cb/26e4a47364bbfdb3b7fb3363be6d8a1c543bcd70a7753ab397350f5f189a/scipy-1.13.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:949ae67db5fa78a86e8fa644b9a6b07252f449dcf74247108c50e1d20d2b4627", size = 33406062, upload-time = "2024-05-23T03:20:36.012Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/88/ab/6ecdc526d509d33814835447bbbeedbebdec7cca46ef495a61b00a35b4bf/scipy-1.13.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:de3ade0e53bc1f21358aa74ff4830235d716211d7d077e340c7349bc3542e884", size = 38218311, upload-time = "2024-05-23T03:20:42.086Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0b/00/9f54554f0f8318100a71515122d8f4f503b1a2c4b4cfab3b4b68c0eb08fa/scipy-1.13.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:2ac65fb503dad64218c228e2dc2d0a0193f7904747db43014645ae139c8fad16", size = 38442493, upload-time = "2024-05-23T03:20:48.292Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3e/df/963384e90733e08eac978cd103c34df181d1fec424de383cdc443f418dd4/scipy-1.13.1-cp312-cp312-win_amd64.whl", hash = "sha256:cdd7dacfb95fea358916410ec61bbc20440f7860333aee6d882bb8046264e949", size = 45910955, upload-time = "2024-05-23T03:20:55.091Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7f/29/c2ea58c9731b9ecb30b6738113a95d147e83922986b34c685b8f6eefde21/scipy-1.13.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:436bbb42a94a8aeef855d755ce5a465479c721e9d684de76bf61a62e7c2b81d5", size = 39352927, upload-time = "2024-05-23T03:21:01.95Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5c/c0/e71b94b20ccf9effb38d7147c0064c08c622309fd487b1b677771a97d18c/scipy-1.13.1-cp39-cp39-macosx_12_0_arm64.whl", hash = "sha256:8335549ebbca860c52bf3d02f80784e91a004b71b059e3eea9678ba994796a24", size = 30324538, upload-time = "2024-05-23T03:21:07.634Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6d/0f/aaa55b06d474817cea311e7b10aab2ea1fd5d43bc6a2861ccc9caec9f418/scipy-1.13.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d533654b7d221a6a97304ab63c41c96473ff04459e404b83275b60aa8f4b7004", size = 33732190, upload-time = "2024-05-23T03:21:14.41Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/35/f5/d0ad1a96f80962ba65e2ce1de6a1e59edecd1f0a7b55990ed208848012e0/scipy-1.13.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:637e98dcf185ba7f8e663e122ebf908c4702420477ae52a04f9908707456ba4d", size = 38612244, upload-time = "2024-05-23T03:21:21.827Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8d/02/1165905f14962174e6569076bcc3315809ae1291ed14de6448cc151eedfd/scipy-1.13.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:a014c2b3697bde71724244f63de2476925596c24285c7a637364761f8710891c", size = 38845637, upload-time = "2024-05-23T03:21:28.729Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3e/77/dab54fe647a08ee4253963bcd8f9cf17509c8ca64d6335141422fe2e2114/scipy-1.13.1-cp39-cp39-win_amd64.whl", hash = "sha256:392e4ec766654852c25ebad4f64e4e584cf19820b980bc04960bca0b0cd6eaa2", size = 46227440, upload-time = "2024-05-23T03:21:35.888Z" }, +] + +[[package]] +name = "scipy" +version = "1.15.3" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'darwin')", +] +dependencies = [ + { name = "numpy", version = "2.2.6", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.10.*'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0f/37/6964b830433e654ec7485e45a00fc9a27cf868d622838f6b6d9c5ec0d532/scipy-1.15.3.tar.gz", hash = "sha256:eae3cf522bc7df64b42cad3925c876e1b0b6c35c1337c93e12c0f366f55b0eaf", size = 59419214, upload-time = "2025-05-08T16:13:05.955Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/78/2f/4966032c5f8cc7e6a60f1b2e0ad686293b9474b65246b0c642e3ef3badd0/scipy-1.15.3-cp310-cp310-macosx_10_13_x86_64.whl", hash = "sha256:a345928c86d535060c9c2b25e71e87c39ab2f22fc96e9636bd74d1dbf9de448c", size = 38702770, upload-time = "2025-05-08T16:04:20.849Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a0/6e/0c3bf90fae0e910c274db43304ebe25a6b391327f3f10b5dcc638c090795/scipy-1.15.3-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:ad3432cb0f9ed87477a8d97f03b763fd1d57709f1bbde3c9369b1dff5503b253", size = 30094511, upload-time = "2025-05-08T16:04:27.103Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ea/b1/4deb37252311c1acff7f101f6453f0440794f51b6eacb1aad4459a134081/scipy-1.15.3-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:aef683a9ae6eb00728a542b796f52a5477b78252edede72b8327a886ab63293f", size = 22368151, upload-time = "2025-05-08T16:04:31.731Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/38/7d/f457626e3cd3c29b3a49ca115a304cebb8cc6f31b04678f03b216899d3c6/scipy-1.15.3-cp310-cp310-macosx_14_0_x86_64.whl", hash = "sha256:1c832e1bd78dea67d5c16f786681b28dd695a8cb1fb90af2e27580d3d0967e92", size = 25121732, upload-time = "2025-05-08T16:04:36.596Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/db/0a/92b1de4a7adc7a15dcf5bddc6e191f6f29ee663b30511ce20467ef9b82e4/scipy-1.15.3-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:263961f658ce2165bbd7b99fa5135195c3a12d9bef045345016b8b50c315cb82", size = 35547617, upload-time = "2025-05-08T16:04:43.546Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8e/6d/41991e503e51fc1134502694c5fa7a1671501a17ffa12716a4a9151af3df/scipy-1.15.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9e2abc762b0811e09a0d3258abee2d98e0c703eee49464ce0069590846f31d40", size = 37662964, upload-time = "2025-05-08T16:04:49.431Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/25/e1/3df8f83cb15f3500478c889be8fb18700813b95e9e087328230b98d547ff/scipy-1.15.3-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:ed7284b21a7a0c8f1b6e5977ac05396c0d008b89e05498c8b7e8f4a1423bba0e", size = 37238749, upload-time = "2025-05-08T16:04:55.215Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/93/3e/b3257cf446f2a3533ed7809757039016b74cd6f38271de91682aa844cfc5/scipy-1.15.3-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:5380741e53df2c566f4d234b100a484b420af85deb39ea35a1cc1be84ff53a5c", size = 40022383, upload-time = "2025-05-08T16:05:01.914Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d1/84/55bc4881973d3f79b479a5a2e2df61c8c9a04fcb986a213ac9c02cfb659b/scipy-1.15.3-cp310-cp310-win_amd64.whl", hash = "sha256:9d61e97b186a57350f6d6fd72640f9e99d5a4a2b8fbf4b9ee9a841eab327dc13", size = 41259201, upload-time = "2025-05-08T16:05:08.166Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/96/ab/5cc9f80f28f6a7dff646c5756e559823614a42b1939d86dd0ed550470210/scipy-1.15.3-cp311-cp311-macosx_10_13_x86_64.whl", hash = "sha256:993439ce220d25e3696d1b23b233dd010169b62f6456488567e830654ee37a6b", size = 38714255, upload-time = "2025-05-08T16:05:14.596Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4a/4a/66ba30abe5ad1a3ad15bfb0b59d22174012e8056ff448cb1644deccbfed2/scipy-1.15.3-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:34716e281f181a02341ddeaad584205bd2fd3c242063bd3423d61ac259ca7eba", size = 30111035, upload-time = "2025-05-08T16:05:20.152Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4b/fa/a7e5b95afd80d24313307f03624acc65801846fa75599034f8ceb9e2cbf6/scipy-1.15.3-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:3b0334816afb8b91dab859281b1b9786934392aa3d527cd847e41bb6f45bee65", size = 22384499, upload-time = "2025-05-08T16:05:24.494Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/17/99/f3aaddccf3588bb4aea70ba35328c204cadd89517a1612ecfda5b2dd9d7a/scipy-1.15.3-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:6db907c7368e3092e24919b5e31c76998b0ce1684d51a90943cb0ed1b4ffd6c1", size = 25152602, upload-time = "2025-05-08T16:05:29.313Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/56/c5/1032cdb565f146109212153339f9cb8b993701e9fe56b1c97699eee12586/scipy-1.15.3-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:721d6b4ef5dc82ca8968c25b111e307083d7ca9091bc38163fb89243e85e3889", size = 35503415, upload-time = "2025-05-08T16:05:34.699Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/bd/37/89f19c8c05505d0601ed5650156e50eb881ae3918786c8fd7262b4ee66d3/scipy-1.15.3-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:39cb9c62e471b1bb3750066ecc3a3f3052b37751c7c3dfd0fd7e48900ed52982", size = 37652622, upload-time = "2025-05-08T16:05:40.762Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7e/31/be59513aa9695519b18e1851bb9e487de66f2d31f835201f1b42f5d4d475/scipy-1.15.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:795c46999bae845966368a3c013e0e00947932d68e235702b5c3f6ea799aa8c9", size = 37244796, upload-time = "2025-05-08T16:05:48.119Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/10/c0/4f5f3eeccc235632aab79b27a74a9130c6c35df358129f7ac8b29f562ac7/scipy-1.15.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:18aaacb735ab38b38db42cb01f6b92a2d0d4b6aabefeb07f02849e47f8fb3594", size = 40047684, upload-time = "2025-05-08T16:05:54.22Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ab/a7/0ddaf514ce8a8714f6ed243a2b391b41dbb65251affe21ee3077ec45ea9a/scipy-1.15.3-cp311-cp311-win_amd64.whl", hash = "sha256:ae48a786a28412d744c62fd7816a4118ef97e5be0bee968ce8f0a2fba7acf3bb", size = 41246504, upload-time = "2025-05-08T16:06:00.437Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/37/4b/683aa044c4162e10ed7a7ea30527f2cbd92e6999c10a8ed8edb253836e9c/scipy-1.15.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:6ac6310fdbfb7aa6612408bd2f07295bcbd3fda00d2d702178434751fe48e019", size = 38766735, upload-time = "2025-05-08T16:06:06.471Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7b/7e/f30be3d03de07f25dc0ec926d1681fed5c732d759ac8f51079708c79e680/scipy-1.15.3-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:185cd3d6d05ca4b44a8f1595af87f9c372bb6acf9c808e99aa3e9aa03bd98cf6", size = 30173284, upload-time = "2025-05-08T16:06:11.686Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/07/9c/0ddb0d0abdabe0d181c1793db51f02cd59e4901da6f9f7848e1f96759f0d/scipy-1.15.3-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:05dc6abcd105e1a29f95eada46d4a3f251743cfd7d3ae8ddb4088047f24ea477", size = 22446958, upload-time = "2025-05-08T16:06:15.97Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/af/43/0bce905a965f36c58ff80d8bea33f1f9351b05fad4beaad4eae34699b7a1/scipy-1.15.3-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:06efcba926324df1696931a57a176c80848ccd67ce6ad020c810736bfd58eb1c", size = 25242454, upload-time = "2025-05-08T16:06:20.394Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/56/30/a6f08f84ee5b7b28b4c597aca4cbe545535c39fe911845a96414700b64ba/scipy-1.15.3-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c05045d8b9bfd807ee1b9f38761993297b10b245f012b11b13b91ba8945f7e45", size = 35210199, upload-time = "2025-05-08T16:06:26.159Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0b/1f/03f52c282437a168ee2c7c14a1a0d0781a9a4a8962d84ac05c06b4c5b555/scipy-1.15.3-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:271e3713e645149ea5ea3e97b57fdab61ce61333f97cfae392c28ba786f9bb49", size = 37309455, upload-time = "2025-05-08T16:06:32.778Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/89/b1/fbb53137f42c4bf630b1ffdfc2151a62d1d1b903b249f030d2b1c0280af8/scipy-1.15.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:6cfd56fc1a8e53f6e89ba3a7a7251f7396412d655bca2aa5611c8ec9a6784a1e", size = 36885140, upload-time = "2025-05-08T16:06:39.249Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2e/2e/025e39e339f5090df1ff266d021892694dbb7e63568edcfe43f892fa381d/scipy-1.15.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:0ff17c0bb1cb32952c09217d8d1eed9b53d1463e5f1dd6052c7857f83127d539", size = 39710549, upload-time = "2025-05-08T16:06:45.729Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e6/eb/3bf6ea8ab7f1503dca3a10df2e4b9c3f6b3316df07f6c0ded94b281c7101/scipy-1.15.3-cp312-cp312-win_amd64.whl", hash = "sha256:52092bc0472cfd17df49ff17e70624345efece4e1a12b23783a1ac59a1b728ed", size = 40966184, upload-time = "2025-05-08T16:06:52.623Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/73/18/ec27848c9baae6e0d6573eda6e01a602e5649ee72c27c3a8aad673ebecfd/scipy-1.15.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:2c620736bcc334782e24d173c0fdbb7590a0a436d2fdf39310a8902505008759", size = 38728256, upload-time = "2025-05-08T16:06:58.696Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/74/cd/1aef2184948728b4b6e21267d53b3339762c285a46a274ebb7863c9e4742/scipy-1.15.3-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:7e11270a000969409d37ed399585ee530b9ef6aa99d50c019de4cb01e8e54e62", size = 30109540, upload-time = "2025-05-08T16:07:04.209Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5b/d8/59e452c0a255ec352bd0a833537a3bc1bfb679944c4938ab375b0a6b3a3e/scipy-1.15.3-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:8c9ed3ba2c8a2ce098163a9bdb26f891746d02136995df25227a20e71c396ebb", size = 22383115, upload-time = "2025-05-08T16:07:08.998Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/08/f5/456f56bbbfccf696263b47095291040655e3cbaf05d063bdc7c7517f32ac/scipy-1.15.3-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:0bdd905264c0c9cfa74a4772cdb2070171790381a5c4d312c973382fc6eaf730", size = 25163884, upload-time = "2025-05-08T16:07:14.091Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a2/66/a9618b6a435a0f0c0b8a6d0a2efb32d4ec5a85f023c2b79d39512040355b/scipy-1.15.3-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:79167bba085c31f38603e11a267d862957cbb3ce018d8b38f79ac043bc92d825", size = 35174018, upload-time = "2025-05-08T16:07:19.427Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b5/09/c5b6734a50ad4882432b6bb7c02baf757f5b2f256041da5df242e2d7e6b6/scipy-1.15.3-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c9deabd6d547aee2c9a81dee6cc96c6d7e9a9b1953f74850c179f91fdc729cb7", size = 37269716, upload-time = "2025-05-08T16:07:25.712Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/77/0a/eac00ff741f23bcabd352731ed9b8995a0a60ef57f5fd788d611d43d69a1/scipy-1.15.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:dde4fc32993071ac0c7dd2d82569e544f0bdaff66269cb475e0f369adad13f11", size = 36872342, upload-time = "2025-05-08T16:07:31.468Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fe/54/4379be86dd74b6ad81551689107360d9a3e18f24d20767a2d5b9253a3f0a/scipy-1.15.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:f77f853d584e72e874d87357ad70f44b437331507d1c311457bed8ed2b956126", size = 39670869, upload-time = "2025-05-08T16:07:38.002Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/87/2e/892ad2862ba54f084ffe8cc4a22667eaf9c2bcec6d2bff1d15713c6c0703/scipy-1.15.3-cp313-cp313-win_amd64.whl", hash = "sha256:b90ab29d0c37ec9bf55424c064312930ca5f4bde15ee8619ee44e69319aab163", size = 40988851, upload-time = "2025-05-08T16:08:33.671Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1b/e9/7a879c137f7e55b30d75d90ce3eb468197646bc7b443ac036ae3fe109055/scipy-1.15.3-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:3ac07623267feb3ae308487c260ac684b32ea35fd81e12845039952f558047b8", size = 38863011, upload-time = "2025-05-08T16:07:44.039Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/51/d1/226a806bbd69f62ce5ef5f3ffadc35286e9fbc802f606a07eb83bf2359de/scipy-1.15.3-cp313-cp313t-macosx_12_0_arm64.whl", hash = "sha256:6487aa99c2a3d509a5227d9a5e889ff05830a06b2ce08ec30df6d79db5fcd5c5", size = 30266407, upload-time = "2025-05-08T16:07:49.891Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e5/9b/f32d1d6093ab9eeabbd839b0f7619c62e46cc4b7b6dbf05b6e615bbd4400/scipy-1.15.3-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:50f9e62461c95d933d5c5ef4a1f2ebf9a2b4e83b0db374cb3f1de104d935922e", size = 22540030, upload-time = "2025-05-08T16:07:54.121Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e7/29/c278f699b095c1a884f29fda126340fcc201461ee8bfea5c8bdb1c7c958b/scipy-1.15.3-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:14ed70039d182f411ffc74789a16df3835e05dc469b898233a245cdfd7f162cb", size = 25218709, upload-time = "2025-05-08T16:07:58.506Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/24/18/9e5374b617aba742a990581373cd6b68a2945d65cc588482749ef2e64467/scipy-1.15.3-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0a769105537aa07a69468a0eefcd121be52006db61cdd8cac8a0e68980bbb723", size = 34809045, upload-time = "2025-05-08T16:08:03.929Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e1/fe/9c4361e7ba2927074360856db6135ef4904d505e9b3afbbcb073c4008328/scipy-1.15.3-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9db984639887e3dffb3928d118145ffe40eff2fa40cb241a306ec57c219ebbbb", size = 36703062, upload-time = "2025-05-08T16:08:09.558Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b7/8e/038ccfe29d272b30086b25a4960f757f97122cb2ec42e62b460d02fe98e9/scipy-1.15.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:40e54d5c7e7ebf1aa596c374c49fa3135f04648a0caabcb66c52884b943f02b4", size = 36393132, upload-time = "2025-05-08T16:08:15.34Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/10/7e/5c12285452970be5bdbe8352c619250b97ebf7917d7a9a9e96b8a8140f17/scipy-1.15.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:5e721fed53187e71d0ccf382b6bf977644c533e506c4d33c3fb24de89f5c3ed5", size = 38979503, upload-time = "2025-05-08T16:08:21.513Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/81/06/0a5e5349474e1cbc5757975b21bd4fad0e72ebf138c5592f191646154e06/scipy-1.15.3-cp313-cp313t-win_amd64.whl", hash = "sha256:76ad1fb5f8752eabf0fa02e4cc0336b4e8f021e2d5f061ed37d6d264db35e3ca", size = 40308097, upload-time = "2025-05-08T16:08:27.627Z" }, +] + +[[package]] +name = "scipy" +version = "1.16.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.13.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'darwin')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'darwin')", + "(python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform == 'darwin')", +] +dependencies = [ + { name = "numpy", version = "2.2.6", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.11'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f5/4a/b927028464795439faec8eaf0b03b011005c487bb2d07409f28bf30879c4/scipy-1.16.1.tar.gz", hash = "sha256:44c76f9e8b6e8e488a586190ab38016e4ed2f8a038af7cd3defa903c0a2238b3", size = 30580861, upload-time = "2025-07-27T16:33:30.834Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/da/91/812adc6f74409b461e3a5fa97f4f74c769016919203138a3bf6fc24ba4c5/scipy-1.16.1-cp311-cp311-macosx_10_14_x86_64.whl", hash = "sha256:c033fa32bab91dc98ca59d0cf23bb876454e2bb02cbe592d5023138778f70030", size = 36552519, upload-time = "2025-07-27T16:26:29.658Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/47/18/8e355edcf3b71418d9e9f9acd2708cc3a6c27e8f98fde0ac34b8a0b45407/scipy-1.16.1-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:6e5c2f74e5df33479b5cd4e97a9104c511518fbd979aa9b8f6aec18b2e9ecae7", size = 28638010, upload-time = "2025-07-27T16:26:38.196Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d9/eb/e931853058607bdfbc11b86df19ae7a08686121c203483f62f1ecae5989c/scipy-1.16.1-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:0a55ffe0ba0f59666e90951971a884d1ff6f4ec3275a48f472cfb64175570f77", size = 20909790, upload-time = "2025-07-27T16:26:43.93Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/45/0c/be83a271d6e96750cd0be2e000f35ff18880a46f05ce8b5d3465dc0f7a2a/scipy-1.16.1-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:f8a5d6cd147acecc2603fbd382fed6c46f474cccfcf69ea32582e033fb54dcfe", size = 23513352, upload-time = "2025-07-27T16:26:50.017Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7c/bf/fe6eb47e74f762f933cca962db7f2c7183acfdc4483bd1c3813cfe83e538/scipy-1.16.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb18899127278058bcc09e7b9966d41a5a43740b5bb8dcba401bd983f82e885b", size = 33534643, upload-time = "2025-07-27T16:26:57.503Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/bb/ba/63f402e74875486b87ec6506a4f93f6d8a0d94d10467280f3d9d7837ce3a/scipy-1.16.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adccd93a2fa937a27aae826d33e3bfa5edf9aa672376a4852d23a7cd67a2e5b7", size = 35376776, upload-time = "2025-07-27T16:27:06.639Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c3/b4/04eb9d39ec26a1b939689102da23d505ea16cdae3dbb18ffc53d1f831044/scipy-1.16.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:18aca1646a29ee9a0625a1be5637fa798d4d81fdf426481f06d69af828f16958", size = 35698906, upload-time = "2025-07-27T16:27:14.943Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/04/d6/bb5468da53321baeb001f6e4e0d9049eadd175a4a497709939128556e3ec/scipy-1.16.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:d85495cef541729a70cdddbbf3e6b903421bc1af3e8e3a9a72a06751f33b7c39", size = 38129275, upload-time = "2025-07-27T16:27:23.873Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c4/94/994369978509f227cba7dfb9e623254d0d5559506fe994aef4bea3ed469c/scipy-1.16.1-cp311-cp311-win_amd64.whl", hash = "sha256:226652fca853008119c03a8ce71ffe1b3f6d2844cc1686e8f9806edafae68596", size = 38644572, upload-time = "2025-07-27T16:27:32.637Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f8/d9/ec4864f5896232133f51382b54a08de91a9d1af7a76dfa372894026dfee2/scipy-1.16.1-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:81b433bbeaf35728dad619afc002db9b189e45eebe2cd676effe1fb93fef2b9c", size = 36575194, upload-time = "2025-07-27T16:27:41.321Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5c/6d/40e81ecfb688e9d25d34a847dca361982a6addf8e31f0957b1a54fbfa994/scipy-1.16.1-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:886cc81fdb4c6903a3bb0464047c25a6d1016fef77bb97949817d0c0d79f9e04", size = 28594590, upload-time = "2025-07-27T16:27:49.204Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0e/37/9f65178edfcc629377ce9a64fc09baebea18c80a9e57ae09a52edf84880b/scipy-1.16.1-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:15240c3aac087a522b4eaedb09f0ad061753c5eebf1ea430859e5bf8640d5919", size = 20866458, upload-time = "2025-07-27T16:27:54.98Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2c/7b/749a66766871ea4cb1d1ea10f27004db63023074c22abed51f22f09770e0/scipy-1.16.1-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:65f81a25805f3659b48126b5053d9e823d3215e4a63730b5e1671852a1705921", size = 23539318, upload-time = "2025-07-27T16:28:01.604Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c4/db/8d4afec60eb833a666434d4541a3151eedbf2494ea6d4d468cbe877f00cd/scipy-1.16.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6c62eea7f607f122069b9bad3f99489ddca1a5173bef8a0c75555d7488b6f725", size = 33292899, upload-time = "2025-07-27T16:28:09.147Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/51/1e/79023ca3bbb13a015d7d2757ecca3b81293c663694c35d6541b4dca53e98/scipy-1.16.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f965bbf3235b01c776115ab18f092a95aa74c271a52577bcb0563e85738fd618", size = 35162637, upload-time = "2025-07-27T16:28:17.535Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b6/49/0648665f9c29fdaca4c679182eb972935b3b4f5ace41d323c32352f29816/scipy-1.16.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:f006e323874ffd0b0b816d8c6a8e7f9a73d55ab3b8c3f72b752b226d0e3ac83d", size = 35490507, upload-time = "2025-07-27T16:28:25.705Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/62/8f/66cbb9d6bbb18d8c658f774904f42a92078707a7c71e5347e8bf2f52bb89/scipy-1.16.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:e8fd15fc5085ab4cca74cb91fe0a4263b1f32e4420761ddae531ad60934c2119", size = 37923998, upload-time = "2025-07-27T16:28:34.339Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/14/c3/61f273ae550fbf1667675701112e380881905e28448c080b23b5a181df7c/scipy-1.16.1-cp312-cp312-win_amd64.whl", hash = "sha256:f7b8013c6c066609577d910d1a2a077021727af07b6fab0ee22c2f901f22352a", size = 38508060, upload-time = "2025-07-27T16:28:43.242Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/93/0b/b5c99382b839854a71ca9482c684e3472badc62620287cbbdab499b75ce6/scipy-1.16.1-cp313-cp313-macosx_10_14_x86_64.whl", hash = "sha256:5451606823a5e73dfa621a89948096c6528e2896e40b39248295d3a0138d594f", size = 36533717, upload-time = "2025-07-27T16:28:51.706Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/eb/e5/69ab2771062c91e23e07c12e7d5033a6b9b80b0903ee709c3c36b3eb520c/scipy-1.16.1-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:89728678c5ca5abd610aee148c199ac1afb16e19844401ca97d43dc548a354eb", size = 28570009, upload-time = "2025-07-27T16:28:57.017Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f4/69/bd75dbfdd3cf524f4d753484d723594aed62cfaac510123e91a6686d520b/scipy-1.16.1-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:e756d688cb03fd07de0fffad475649b03cb89bee696c98ce508b17c11a03f95c", size = 20841942, upload-time = "2025-07-27T16:29:01.152Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ea/74/add181c87663f178ba7d6144b370243a87af8476664d5435e57d599e6874/scipy-1.16.1-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:5aa2687b9935da3ed89c5dbed5234576589dd28d0bf7cd237501ccfbdf1ad608", size = 23498507, upload-time = "2025-07-27T16:29:05.202Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1d/74/ece2e582a0d9550cee33e2e416cc96737dce423a994d12bbe59716f47ff1/scipy-1.16.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:0851f6a1e537fe9399f35986897e395a1aa61c574b178c0d456be5b1a0f5ca1f", size = 33286040, upload-time = "2025-07-27T16:29:10.201Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e4/82/08e4076df538fb56caa1d489588d880ec7c52d8273a606bb54d660528f7c/scipy-1.16.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fedc2cbd1baed37474b1924c331b97bdff611d762c196fac1a9b71e67b813b1b", size = 35176096, upload-time = "2025-07-27T16:29:17.091Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fa/79/cd710aab8c921375711a8321c6be696e705a120e3011a643efbbcdeeabcc/scipy-1.16.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:2ef500e72f9623a6735769e4b93e9dcb158d40752cdbb077f305487e3e2d1f45", size = 35490328, upload-time = "2025-07-27T16:29:22.928Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/71/73/e9cc3d35ee4526d784520d4494a3e1ca969b071fb5ae5910c036a375ceec/scipy-1.16.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:978d8311674b05a8f7ff2ea6c6bce5d8b45a0cb09d4c5793e0318f448613ea65", size = 37939921, upload-time = "2025-07-27T16:29:29.108Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/21/12/c0efd2941f01940119b5305c375ae5c0fcb7ec193f806bd8f158b73a1782/scipy-1.16.1-cp313-cp313-win_amd64.whl", hash = "sha256:81929ed0fa7a5713fcdd8b2e6f73697d3b4c4816d090dd34ff937c20fa90e8ab", size = 38479462, upload-time = "2025-07-27T16:30:24.078Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7a/19/c3d08b675260046a991040e1ea5d65f91f40c7df1045fffff412dcfc6765/scipy-1.16.1-cp313-cp313t-macosx_10_14_x86_64.whl", hash = "sha256:bcc12db731858abda693cecdb3bdc9e6d4bd200213f49d224fe22df82687bdd6", size = 36938832, upload-time = "2025-07-27T16:29:35.057Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/81/f2/ce53db652c033a414a5b34598dba6b95f3d38153a2417c5a3883da429029/scipy-1.16.1-cp313-cp313t-macosx_12_0_arm64.whl", hash = "sha256:744d977daa4becb9fc59135e75c069f8d301a87d64f88f1e602a9ecf51e77b27", size = 29093084, upload-time = "2025-07-27T16:29:40.201Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a9/ae/7a10ff04a7dc15f9057d05b33737ade244e4bd195caa3f7cc04d77b9e214/scipy-1.16.1-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:dc54f76ac18073bcecffb98d93f03ed6b81a92ef91b5d3b135dcc81d55a724c7", size = 21365098, upload-time = "2025-07-27T16:29:44.295Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/36/ac/029ff710959932ad3c2a98721b20b405f05f752f07344622fd61a47c5197/scipy-1.16.1-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:367d567ee9fc1e9e2047d31f39d9d6a7a04e0710c86e701e053f237d14a9b4f6", size = 23896858, upload-time = "2025-07-27T16:29:48.784Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/71/13/d1ef77b6bd7898720e1f0b6b3743cb945f6c3cafa7718eaac8841035ab60/scipy-1.16.1-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:4cf5785e44e19dcd32a0e4807555e1e9a9b8d475c6afff3d21c3c543a6aa84f4", size = 33438311, upload-time = "2025-07-27T16:29:54.164Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2d/e0/e64a6821ffbb00b4c5b05169f1c1fddb4800e9307efe3db3788995a82a2c/scipy-1.16.1-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3d0b80fb26d3e13a794c71d4b837e2a589d839fd574a6bbb4ee1288c213ad4a3", size = 35279542, upload-time = "2025-07-27T16:30:00.249Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/57/59/0dc3c8b43e118f1e4ee2b798dcc96ac21bb20014e5f1f7a8e85cc0653bdb/scipy-1.16.1-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:8503517c44c18d1030d666cb70aaac1cc8913608816e06742498833b128488b7", size = 35667665, upload-time = "2025-07-27T16:30:05.916Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/45/5f/844ee26e34e2f3f9f8febb9343748e72daeaec64fe0c70e9bf1ff84ec955/scipy-1.16.1-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:30cc4bb81c41831ecfd6dc450baf48ffd80ef5aed0f5cf3ea775740e80f16ecc", size = 38045210, upload-time = "2025-07-27T16:30:11.655Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8d/d7/210f2b45290f444f1de64bc7353aa598ece9f0e90c384b4a156f9b1a5063/scipy-1.16.1-cp313-cp313t-win_amd64.whl", hash = "sha256:c24fa02f7ed23ae514460a22c57eca8f530dbfa50b1cfdbf4f37c05b5309cc39", size = 38593661, upload-time = "2025-07-27T16:30:17.825Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/81/ea/84d481a5237ed223bd3d32d6e82d7a6a96e34756492666c260cef16011d1/scipy-1.16.1-cp314-cp314-macosx_10_14_x86_64.whl", hash = "sha256:796a5a9ad36fa3a782375db8f4241ab02a091308eb079746bc0f874c9b998318", size = 36525921, upload-time = "2025-07-27T16:30:30.081Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4e/9f/d9edbdeff9f3a664807ae3aea383e10afaa247e8e6255e6d2aa4515e8863/scipy-1.16.1-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:3ea0733a2ff73fd6fdc5fecca54ee9b459f4d74f00b99aced7d9a3adb43fb1cc", size = 28564152, upload-time = "2025-07-27T16:30:35.336Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3b/95/8125bcb1fe04bc267d103e76516243e8d5e11229e6b306bda1024a5423d1/scipy-1.16.1-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:85764fb15a2ad994e708258bb4ed8290d1305c62a4e1ef07c414356a24fcfbf8", size = 20836028, upload-time = "2025-07-27T16:30:39.421Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/77/9c/bf92e215701fc70bbcd3d14d86337cf56a9b912a804b9c776a269524a9e9/scipy-1.16.1-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:ca66d980469cb623b1759bdd6e9fd97d4e33a9fad5b33771ced24d0cb24df67e", size = 23489666, upload-time = "2025-07-27T16:30:43.663Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5e/00/5e941d397d9adac41b02839011594620d54d99488d1be5be755c00cde9ee/scipy-1.16.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e7cc1ffcc230f568549fc56670bcf3df1884c30bd652c5da8138199c8c76dae0", size = 33358318, upload-time = "2025-07-27T16:30:48.982Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0e/87/8db3aa10dde6e3e8e7eb0133f24baa011377d543f5b19c71469cf2648026/scipy-1.16.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3ddfb1e8d0b540cb4ee9c53fc3dea3186f97711248fb94b4142a1b27178d8b4b", size = 35185724, upload-time = "2025-07-27T16:30:54.26Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/89/b4/6ab9ae443216807622bcff02690262d8184078ea467efee2f8c93288a3b1/scipy-1.16.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:4dc0e7be79e95d8ba3435d193e0d8ce372f47f774cffd882f88ea4e1e1ddc731", size = 35554335, upload-time = "2025-07-27T16:30:59.765Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/9c/9a/d0e9dc03c5269a1afb60661118296a32ed5d2c24298af61b676c11e05e56/scipy-1.16.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:f23634f9e5adb51b2a77766dac217063e764337fbc816aa8ad9aaebcd4397fd3", size = 37960310, upload-time = "2025-07-27T16:31:06.151Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5e/00/c8f3130a50521a7977874817ca89e0599b1b4ee8e938bad8ae798a0e1f0d/scipy-1.16.1-cp314-cp314-win_amd64.whl", hash = "sha256:57d75524cb1c5a374958a2eae3d84e1929bb971204cc9d52213fb8589183fc19", size = 39319239, upload-time = "2025-07-27T16:31:59.942Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f2/f2/1ca3eda54c3a7e4c92f6acef7db7b3a057deb135540d23aa6343ef8ad333/scipy-1.16.1-cp314-cp314t-macosx_10_14_x86_64.whl", hash = "sha256:d8da7c3dd67bcd93f15618938f43ed0995982eb38973023d46d4646c4283ad65", size = 36939460, upload-time = "2025-07-27T16:31:11.865Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/80/30/98c2840b293a132400c0940bb9e140171dcb8189588619048f42b2ce7b4f/scipy-1.16.1-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:cc1d2f2fd48ba1e0620554fe5bc44d3e8f5d4185c8c109c7fbdf5af2792cfad2", size = 29093322, upload-time = "2025-07-27T16:31:17.045Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c1/e6/1e6e006e850622cf2a039b62d1a6ddc4497d4851e58b68008526f04a9a00/scipy-1.16.1-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:21a611ced9275cb861bacadbada0b8c0623bc00b05b09eb97f23b370fc2ae56d", size = 21365329, upload-time = "2025-07-27T16:31:21.188Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8e/02/72a5aa5b820589dda9a25e329ca752842bfbbaf635e36bc7065a9b42216e/scipy-1.16.1-cp314-cp314t-macosx_14_0_x86_64.whl", hash = "sha256:8dfbb25dffc4c3dd9371d8ab456ca81beeaf6f9e1c2119f179392f0dc1ab7695", size = 23897544, upload-time = "2025-07-27T16:31:25.408Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2b/dc/7122d806a6f9eb8a33532982234bed91f90272e990f414f2830cfe656e0b/scipy-1.16.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f0ebb7204f063fad87fc0a0e4ff4a2ff40b2a226e4ba1b7e34bf4b79bf97cd86", size = 33442112, upload-time = "2025-07-27T16:31:30.62Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/24/39/e383af23564daa1021a5b3afbe0d8d6a68ec639b943661841f44ac92de85/scipy-1.16.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f1b9e5962656f2734c2b285a8745358ecb4e4efbadd00208c80a389227ec61ff", size = 35286594, upload-time = "2025-07-27T16:31:36.112Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/95/47/1a0b0aff40c3056d955f38b0df5d178350c3d74734ec54f9c68d23910be5/scipy-1.16.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:5e1a106f8c023d57a2a903e771228bf5c5b27b5d692088f457acacd3b54511e4", size = 35665080, upload-time = "2025-07-27T16:31:42.025Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/64/df/ce88803e9ed6e27fe9b9abefa157cf2c80e4fa527cf17ee14be41f790ad4/scipy-1.16.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:709559a1db68a9abc3b2c8672c4badf1614f3b440b3ab326d86a5c0491eafae3", size = 38050306, upload-time = "2025-07-27T16:31:48.109Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6e/6c/a76329897a7cae4937d403e623aa6aaea616a0bb5b36588f0b9d1c9a3739/scipy-1.16.1-cp314-cp314t-win_amd64.whl", hash = "sha256:c0c804d60492a0aad7f5b2bb1862f4548b990049e27e828391ff2bf6f7199998", size = 39427705, upload-time = "2025-07-27T16:31:53.96Z" }, +] + +[[package]] +name = "setuptools" +version = "75.3.2" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5c/01/771ea46cce201dd42cff043a5eea929d1c030fb3d1c2ee2729d02ca7814c/setuptools-75.3.2.tar.gz", hash = "sha256:3c1383e1038b68556a382c1e8ded8887cd20141b0eb5708a6c8d277de49364f5", size = 1354489, upload-time = "2025-03-12T00:02:19.004Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/15/65/3f0dba35760d902849d39d38c0a72767794b1963227b69a587f8a336d08c/setuptools-75.3.2-py3-none-any.whl", hash = "sha256:90ab613b6583fc02d5369cbca13ea26ea0e182d1df2d943ee9cbe81d4c61add9", size = 1251198, upload-time = "2025-03-12T00:02:17.554Z" }, +] + +[[package]] +name = "setuptools" +version = "80.9.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.13.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'darwin')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'darwin')", + "(python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform == 'darwin')", + "(python_full_version == '3.10.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'darwin')", + "(python_full_version == '3.9.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform == 'darwin')", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/18/5d/3bf57dcd21979b887f014ea83c24ae194cfcd12b9e0fda66b957c69d1fca/setuptools-80.9.0.tar.gz", hash = "sha256:f36b47402ecde768dbfafc46e8e4207b4360c654f1f3bb84475f0a28628fb19c", size = 1319958, upload-time = "2025-05-27T00:56:51.443Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a3/dc/17031897dae0efacfea57dfd3a82fdd2a2aeb58e0ff71b77b87e44edc772/setuptools-80.9.0-py3-none-any.whl", hash = "sha256:062d34222ad13e0cc312a4c02d73f059e86a4acbfbdea8f8f76b28c99f306922", size = 1201486, upload-time = "2025-05-27T00:56:49.664Z" }, +] + +[[package]] +name = "six" +version = "1.17.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/94/e7/b2c673351809dca68a0e064b6af791aa332cf192da575fd474ed7d6f16a2/six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81", size = 34031, upload-time = "2024-12-04T17:35:28.174Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, +] + +[[package]] +name = "soundfile" +version = "0.13.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +dependencies = [ + { name = "cffi" }, + { name = "numpy", version = "1.24.4", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "numpy", version = "2.0.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.10'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e1/41/9b873a8c055582859b239be17902a85339bec6a30ad162f98c9b0288a2cc/soundfile-0.13.1.tar.gz", hash = "sha256:b2c68dab1e30297317080a5b43df57e302584c49e2942defdde0acccc53f0e5b", size = 46156, upload-time = "2025-01-25T09:17:04.831Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/64/28/e2a36573ccbcf3d57c00626a21fe51989380636e821b341d36ccca0c1c3a/soundfile-0.13.1-py2.py3-none-any.whl", hash = "sha256:a23c717560da2cf4c7b5ae1142514e0fd82d6bbd9dfc93a50423447142f2c445", size = 25751, upload-time = "2025-01-25T09:16:44.235Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ea/ab/73e97a5b3cc46bba7ff8650a1504348fa1863a6f9d57d7001c6b67c5f20e/soundfile-0.13.1-py2.py3-none-macosx_10_9_x86_64.whl", hash = "sha256:82dc664d19831933fe59adad199bf3945ad06d84bc111a5b4c0d3089a5b9ec33", size = 1142250, upload-time = "2025-01-25T09:16:47.583Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a0/e5/58fd1a8d7b26fc113af244f966ee3aecf03cb9293cb935daaddc1e455e18/soundfile-0.13.1-py2.py3-none-macosx_11_0_arm64.whl", hash = "sha256:743f12c12c4054921e15736c6be09ac26b3b3d603aef6fd69f9dde68748f2593", size = 1101406, upload-time = "2025-01-25T09:16:49.662Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/58/ae/c0e4a53d77cf6e9a04179535766b3321b0b9ced5f70522e4caf9329f0046/soundfile-0.13.1-py2.py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:9c9e855f5a4d06ce4213f31918653ab7de0c5a8d8107cd2427e44b42df547deb", size = 1235729, upload-time = "2025-01-25T09:16:53.018Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/57/5e/70bdd9579b35003a489fc850b5047beeda26328053ebadc1fb60f320f7db/soundfile-0.13.1-py2.py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:03267c4e493315294834a0870f31dbb3b28a95561b80b134f0bd3cf2d5f0e618", size = 1313646, upload-time = "2025-01-25T09:16:54.872Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fe/df/8c11dc4dfceda14e3003bb81a0d0edcaaf0796dd7b4f826ea3e532146bba/soundfile-0.13.1-py2.py3-none-win32.whl", hash = "sha256:c734564fab7c5ddf8e9be5bf70bab68042cd17e9c214c06e365e20d64f9a69d5", size = 899881, upload-time = "2025-01-25T09:16:56.663Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/14/e9/6b761de83277f2f02ded7e7ea6f07828ec78e4b229b80e4ca55dd205b9dc/soundfile-0.13.1-py2.py3-none-win_amd64.whl", hash = "sha256:1e70a05a0626524a69e9f0f4dd2ec174b4e9567f4d8b6c11d38b5c289be36ee9", size = 1019162, upload-time = "2025-01-25T09:16:59.573Z" }, +] + +[[package]] +name = "soxr" +version = "0.3.7" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] +dependencies = [ + { name = "numpy", version = "1.24.4", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/49/31/f24193672d07964cbf85f2e7af3b45936e1c76e44eb5ad97199865e2dcf8/soxr-0.3.7.tar.gz", hash = "sha256:436ddff00c6eb2c75b79c19cfdca7527b1e31b5fad738652f044045ba6258593", size = 296432, upload-time = "2023-10-05T02:36:28.674Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/22/32/c8f8c60d91db642ee2ecf86656a35bd5463b12b54cc72f3161750e2539d3/soxr-0.3.7-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:ac81c4af6a993d5b7c0b466bbac4835bad2b14ec32f342b2c1f83e4cf825e301", size = 414127, upload-time = "2023-10-05T02:35:44.432Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/bc/38/2635bcf180de54457d64a6b348b3e421f469aee7edafead2306a6e74cc1a/soxr-0.3.7-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:8d8a2b3e7f8d0255e2484fb82cb66c86da6fb25b342ef793cceca9ce9a61aa16", size = 390040, upload-time = "2023-10-05T02:35:46.164Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c0/2a/1c492e2ffa0cddabef105332d488e647b8d1315a21cb881e3671a2812880/soxr-0.3.7-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a0cd6eb6f6bbda2e8de36672cf2f0529ced6e638773150744ef075be0cc4f52c", size = 1216295, upload-time = "2023-10-05T02:35:47.847Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/31/f7/d95b816c47dca6a068305fb7176b8c8d2c94bbc6cce6dcc296c6cf98660f/soxr-0.3.7-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e47d86af35b942c92606fc2d5dfccf3f01309329475571ae2312bbf9edc3a790", size = 1238709, upload-time = "2023-10-05T02:35:49.66Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3c/e7/89951b917600d02f7389e760696c32b70a80a96301d0a018a70a317e4ecc/soxr-0.3.7-cp310-cp310-win_amd64.whl", hash = "sha256:0e291adfaf9f2a7c4dd180a1b8c280f9beb1c84cb381853e4f4b3434d002ed7f", size = 184591, upload-time = "2023-10-05T02:35:51.441Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1c/5a/61572e3d1f5dd4c9c0a32d9f7f57b7b6089222227c881fe0945cf955261f/soxr-0.3.7-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:9e811450f0e91972932bd37ac58e32e44002c2c99db2aa926a9e7ba164545034", size = 414444, upload-time = "2023-10-05T02:35:52.847Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d0/79/c93daceac24cd6830333d9f3d04716115240ef3be4e8d8ece511afbcf417/soxr-0.3.7-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:9cea63014ce91035074e1228c9340e2b8609faf964e268705fcac5135d05060c", size = 390460, upload-time = "2023-10-05T02:35:54.592Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ad/d7/6ec56e463e74764a9b9c2cc2645788991d7ce6c82cd40ff002e5fd6694bc/soxr-0.3.7-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bfab27830f6217a15b83445988225c3aeea3bbccfa9399ced291e53e1b05925d", size = 1287621, upload-time = "2023-10-05T02:35:56.017Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2b/97/cbce72f9c8b5c9c667eb55dc55be20a87c610dba55c0466c77498c1a8c97/soxr-0.3.7-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:286858e3078d76c11b6d490b66fed3c9bb2a4229759f6be03ceef5c02189bf2c", size = 1298574, upload-time = "2023-10-05T02:35:57.819Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0d/95/8398287a3a00d006ebbe5fcada88bb2536acde2beecce892ff091e1c1a3b/soxr-0.3.7-cp311-cp311-win_amd64.whl", hash = "sha256:54985ff33292192d2937be80df3e5f3a44d6d53e6835f727d6b99b7cdd3f1611", size = 184720, upload-time = "2023-10-05T02:35:59.028Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/04/2a/dfb7e28335e9715d0494a7f896f4a5ae70e42761ca7740392dd52fa8c133/soxr-0.3.7-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:83c74ef6d61d7dcd81be26f91bee0a420f792f5c1982266f2a80e655f0650a98", size = 413863, upload-time = "2023-10-05T02:36:00.142Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d6/5b/14ff65df4ba5663b69ddc6e80bc10c0c98f67c3cfd9a3aa435ab8bde3b92/soxr-0.3.7-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:cb1e14663a43fe88b8fbc287822a159028366a820abe1a0a9670fb53618cb47b", size = 390167, upload-time = "2023-10-05T02:36:01.935Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7a/1a/1287a8d6a470d1ce5084269ec7ac17dfc8eeb3e48bec3bd0be0eefed8ed4/soxr-0.3.7-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:48acdfbcf870ab54f645b1cfd641bce92c1e3a67346c3bf0f6c0ad2873c1dd35", size = 1270769, upload-time = "2023-10-05T02:36:03.351Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1b/51/54a6bbb3c0e80b7f59a7e30a4643fef18f43b75ce30512a6baf4afb8d490/soxr-0.3.7-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ea663b76f2b0ec1576b8a43aef317aec080abc0a67a4015fcd9f3407039f260a", size = 1298776, upload-time = "2023-10-05T02:36:05.375Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/86/33/15edab5cd5299238c2b5c67031472108374b8fbd4b488843919cc7a03698/soxr-0.3.7-cp312-cp312-win_amd64.whl", hash = "sha256:42da0d9eb79c70e5a41917f1b48a032e241a48eb4a1bcea7c80577302ff26974", size = 184068, upload-time = "2023-10-05T02:36:06.897Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/39/69/b1cd2acc3642873b836e12197343e52cd5626d0820441092438c6890901a/soxr-0.3.7-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:ce12b93747958f2769d6b297e6e27c73d9ad635fe8104ef052bece9c8a322824", size = 413051, upload-time = "2023-10-05T02:36:14.692Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d3/3d/04a8be9e776bd7c39634066927f3eeed5134e3c3f6c1e6211ce6b69e6025/soxr-0.3.7-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:1cd65dc7b96ea3cb6c8c48e6020e859680556cc42dd3d4de44779530cce21037", size = 389443, upload-time = "2023-10-05T02:36:16.511Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/70/90/872f5a4277fc43ca5643cd9919d8de8d3c392fdd03fc09861669e6c1f2ff/soxr-0.3.7-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d994f1a7690b1b13ab639ea33e0c1d78415b64d88d6df4af705a9443f97b9687", size = 1246233, upload-time = "2023-10-05T02:36:17.89Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ce/6b/f9ac173f4365d52d247df1e4727c932d241b4e8b04f0820b1d5c44ba18c8/soxr-0.3.7-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e87b58bc9e8c2caa16f07726f666bd043f0a49ca937baa803ce7708003b27833", size = 1269620, upload-time = "2023-10-05T02:36:19.052Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e2/8f/bd1c42ae58a41a7f30f908659b510d27b30657312e64d8005f234450d529/soxr-0.3.7-cp38-cp38-win_amd64.whl", hash = "sha256:07f4c0c6125ea1482fa187ad5f007216712ee0a93586a9b2f80e79c0bf944cf7", size = 184966, upload-time = "2023-10-05T02:36:20.312Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/bb/a0/68de2c18fd5079f979ece27b40b86f388c3083f491386d47731c610f3341/soxr-0.3.7-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:e5267c3ba34d4b873d9bbe3a9e58418b01ae4fd04349a4f944d9943b9ddac0f7", size = 413987, upload-time = "2023-10-05T02:36:21.957Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3e/86/7fd03c9f2a825a1434ba1a5097e793fdb431571beaefcd7030d013dca707/soxr-0.3.7-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:6e39668c250e221db888cf3b290a16fbe10a702d9a4eb604a127f720040de583", size = 390064, upload-time = "2023-10-05T02:36:23.045Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ab/8d/d228996ae2d35944f7f79792e0a45a7446e59b4b09367380db5a006cdfd8/soxr-0.3.7-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f8ceeb74e5a55d903cc286d3bd12c2d8f8c85d02894071e9ec92ab405430907c", size = 1220544, upload-time = "2023-10-05T02:36:24.235Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/be/1a/845d95266ea63ff402121ac5d1d48261bebfd20f45dbd24ce08fb712bf47/soxr-0.3.7-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0eed6bf58192dd1bb93becd2444de4d712689713d727b32fd55623ae9aae7df7", size = 1239697, upload-time = "2023-10-05T02:36:25.494Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8b/a1/f1a343ff7d37b00931f76f629ea121df6b639d80765d67b41faa521ce692/soxr-0.3.7-cp39-cp39-win_amd64.whl", hash = "sha256:7221302b4547d02a3f38dd3cd15317ab2b78873c75921db5f4a070848f0c71be", size = 184793, upload-time = "2023-10-05T02:36:26.983Z" }, +] + +[[package]] +name = "soxr" +version = "0.5.0.post1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.13.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'darwin')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'darwin')", + "(python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform == 'darwin')", + "(python_full_version == '3.10.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'darwin')", + "(python_full_version == '3.9.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform == 'darwin')", +] +dependencies = [ + { name = "numpy", version = "2.0.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.10'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/02/c0/4429bf9b3be10e749149e286aa5c53775399ec62891c6b970456c6dca325/soxr-0.5.0.post1.tar.gz", hash = "sha256:7092b9f3e8a416044e1fa138c8172520757179763b85dc53aa9504f4813cff73", size = 170853, upload-time = "2024-08-31T03:43:33.058Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7d/96/bee1eb69d66fc28c3b219ba9b8674b49d3dcc6cd2f9b3e5114ff28cf88b5/soxr-0.5.0.post1-cp310-cp310-macosx_10_14_x86_64.whl", hash = "sha256:7406d782d85f8cf64e66b65e6b7721973de8a1dc50b9e88bc2288c343a987484", size = 203841, upload-time = "2024-08-31T03:42:59.186Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1f/5d/56ad3d181d30d103128f65cc44f4c4e24c199e6d5723e562704e47c89f78/soxr-0.5.0.post1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:fa0a382fb8d8e2afed2c1642723b2d2d1b9a6728ff89f77f3524034c8885b8c9", size = 160192, upload-time = "2024-08-31T03:43:01.128Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7f/09/e43c39390e26b4c1b8d46f8a1c252a5077fa9f81cc2326b03c3d2b85744e/soxr-0.5.0.post1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8b01d3efb95a2851f78414bcd00738b0253eec3f5a1e5482838e965ffef84969", size = 221176, upload-time = "2024-08-31T03:43:02.663Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ba/e6/059070b4cdb7fdd8ffbb67c5087c1da9716577127fb0540cd11dbf77923b/soxr-0.5.0.post1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fcc049b0a151a65aa75b92f0ac64bb2dba785d16b78c31c2b94e68c141751d6d", size = 252779, upload-time = "2024-08-31T03:43:04.582Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ad/64/86082b6372e5ff807dfa79b857da9f50e94e155706000daa43fdc3b59851/soxr-0.5.0.post1-cp310-cp310-win_amd64.whl", hash = "sha256:97f269bc26937c267a2ace43a77167d0c5c8bba5a2b45863bb6042b5b50c474e", size = 166881, upload-time = "2024-08-31T03:43:06.255Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/29/28/dc62dae260a77603e8257e9b79078baa2ca4c0b4edc6f9f82c9113d6ef18/soxr-0.5.0.post1-cp311-cp311-macosx_10_14_x86_64.whl", hash = "sha256:6fb77b626773a966e3d8f6cb24f6f74b5327fa5dc90f1ff492450e9cdc03a378", size = 203648, upload-time = "2024-08-31T03:43:08.339Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0e/48/3e88329a695f6e0e38a3b171fff819d75d7cc055dae1ec5d5074f34d61e3/soxr-0.5.0.post1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:39e0f791ba178d69cd676485dbee37e75a34f20daa478d90341ecb7f6d9d690f", size = 159933, upload-time = "2024-08-31T03:43:10.053Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/9c/a5/6b439164be6871520f3d199554568a7656e96a867adbbe5bac179caf5776/soxr-0.5.0.post1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4f0b558f445ba4b64dbcb37b5f803052eee7d93b1dbbbb97b3ec1787cb5a28eb", size = 221010, upload-time = "2024-08-31T03:43:11.839Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/9f/e5/400e3bf7f29971abad85cb877e290060e5ec61fccd2fa319e3d85709c1be/soxr-0.5.0.post1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ca6903671808e0a6078b0d146bb7a2952b118dfba44008b2aa60f221938ba829", size = 252471, upload-time = "2024-08-31T03:43:13.347Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/86/94/6a7e91bea7e6ca193ee429869b8f18548cd79759e064021ecb5756024c7c/soxr-0.5.0.post1-cp311-cp311-win_amd64.whl", hash = "sha256:c4d8d5283ed6f5efead0df2c05ae82c169cfdfcf5a82999c2d629c78b33775e8", size = 166723, upload-time = "2024-08-31T03:43:15.212Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5d/e3/d422d279e51e6932e7b64f1170a4f61a7ee768e0f84c9233a5b62cd2c832/soxr-0.5.0.post1-cp312-abi3-macosx_10_14_x86_64.whl", hash = "sha256:fef509466c9c25f65eae0ce1e4b9ac9705d22c6038c914160ddaf459589c6e31", size = 199993, upload-time = "2024-08-31T03:43:17.24Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/20/f1/88adaca3c52e03bcb66b63d295df2e2d35bf355d19598c6ce84b20be7fca/soxr-0.5.0.post1-cp312-abi3-macosx_11_0_arm64.whl", hash = "sha256:4704ba6b13a3f1e41d12acf192878384c1c31f71ce606829c64abdf64a8d7d32", size = 156373, upload-time = "2024-08-31T03:43:18.633Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b8/38/bad15a9e615215c8219652ca554b601663ac3b7ac82a284aca53ec2ff48c/soxr-0.5.0.post1-cp312-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bd052a66471a7335b22a6208601a9d0df7b46b8d087dce4ff6e13eed6a33a2a1", size = 216564, upload-time = "2024-08-31T03:43:20.789Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e1/1a/569ea0420a0c4801c2c8dd40d8d544989522f6014d51def689125f3f2935/soxr-0.5.0.post1-cp312-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a3f16810dd649ab1f433991d2a9661e9e6a116c2b4101039b53b3c3e90a094fc", size = 248455, upload-time = "2024-08-31T03:43:22.165Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/bc/10/440f1ba3d4955e0dc740bbe4ce8968c254a3d644d013eb75eea729becdb8/soxr-0.5.0.post1-cp312-abi3-win_amd64.whl", hash = "sha256:b1be9fee90afb38546bdbd7bde714d1d9a8c5a45137f97478a83b65e7f3146f6", size = 164937, upload-time = "2024-08-31T03:43:23.671Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d9/7b/c8d797235d06ae316e0c9bc2b1d0d5d948834dfac17eba0ad10fd177524b/soxr-0.5.0.post1-cp39-cp39-macosx_10_14_x86_64.whl", hash = "sha256:c5af7b355959061beb90a1d73c4834ece4549f07b708f8c73c088153cec29935", size = 204073, upload-time = "2024-08-31T03:43:24.935Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/88/5c/f6cf6b90ce1628def17c746d6cde9991fdd29667ef1d5fb5bd3b22eb788f/soxr-0.5.0.post1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:e1dda616fc797b1507b65486f3116ed2c929f13c722922963dd419d64ada6c07", size = 160375, upload-time = "2024-08-31T03:43:26.742Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a5/4a/6a11d62cfd6383c88f4918bdc5191d9c437f649c9101ceb5eec7e2837f0b/soxr-0.5.0.post1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:94de2812368e98cb42b4eaeddf8ee1657ecc19bd053f8e67b9b5aa12a3592012", size = 221450, upload-time = "2024-08-31T03:43:28.048Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/25/d1/83a66e795381ddfc5c3ebf34cc0ac68735c7c459ed1fe65a2193a52c57b1/soxr-0.5.0.post1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9c8e9c980637e03d3f345a4fd81d56477a58c294fb26205fa121bc4eb23d9d01", size = 253025, upload-time = "2024-08-31T03:43:29.893Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/cd/2e/1fbad5cc8c49c45a0b94221d16445792f55b63fe4f6d3885db960d92892c/soxr-0.5.0.post1-cp39-cp39-win_amd64.whl", hash = "sha256:7e71b0b0db450f36de70f1047505231db77a713f8c47df9342582ae8a4b828f2", size = 167292, upload-time = "2024-08-31T03:43:31.313Z" }, +] + +[[package]] +name = "standard-aifc" +version = "3.13.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +dependencies = [ + { name = "audioop-lts", marker = "python_full_version >= '3.13'" }, + { name = "standard-chunk", marker = "python_full_version >= '3.13'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c4/53/6050dc3dde1671eb3db592c13b55a8005e5040131f7509cef0215212cb84/standard_aifc-3.13.0.tar.gz", hash = "sha256:64e249c7cb4b3daf2fdba4e95721f811bde8bdfc43ad9f936589b7bb2fae2e43", size = 15240, upload-time = "2024-10-30T16:01:31.772Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c3/52/5fbb203394cc852334d1575cc020f6bcec768d2265355984dfd361968f36/standard_aifc-3.13.0-py3-none-any.whl", hash = "sha256:f7ae09cc57de1224a0dd8e3eb8f73830be7c3d0bc485de4c1f82b4a7f645ac66", size = 10492, upload-time = "2024-10-30T16:01:07.071Z" }, +] + +[[package]] +name = "standard-chunk" +version = "3.13.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/43/06/ce1bb165c1f111c7d23a1ad17204d67224baa69725bb6857a264db61beaf/standard_chunk-3.13.0.tar.gz", hash = "sha256:4ac345d37d7e686d2755e01836b8d98eda0d1a3ee90375e597ae43aaf064d654", size = 4672, upload-time = "2024-10-30T16:18:28.326Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7a/90/a5c1084d87767d787a6caba615aa50dc587229646308d9420c960cb5e4c0/standard_chunk-3.13.0-py3-none-any.whl", hash = "sha256:17880a26c285189c644bd5bd8f8ed2bdb795d216e3293e6dbe55bbd848e2982c", size = 4944, upload-time = "2024-10-30T16:18:26.694Z" }, +] + +[[package]] +name = "standard-sunau" +version = "3.13.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +dependencies = [ + { name = "audioop-lts", marker = "python_full_version >= '3.13'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/66/e3/ce8d38cb2d70e05ffeddc28bb09bad77cfef979eb0a299c9117f7ed4e6a9/standard_sunau-3.13.0.tar.gz", hash = "sha256:b319a1ac95a09a2378a8442f403c66f4fd4b36616d6df6ae82b8e536ee790908", size = 9368, upload-time = "2024-10-30T16:01:41.626Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/34/ae/e3707f6c1bc6f7aa0df600ba8075bfb8a19252140cd595335be60e25f9ee/standard_sunau-3.13.0-py3-none-any.whl", hash = "sha256:53af624a9529c41062f4c2fd33837f297f3baa196b0cfceffea6555654602622", size = 7364, upload-time = "2024-10-30T16:01:28.003Z" }, +] + +[[package]] +name = "sympy" +version = "1.12.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] +dependencies = [ + { name = "mpmath", marker = "python_full_version < '3.9'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/41/8a/0d1bbd33cd3091c913d298746e56f40586fa954788f51b816c6336424675/sympy-1.12.1.tar.gz", hash = "sha256:2877b03f998cd8c08f07cd0de5b767119cd3ef40d09f41c30d722f6686b0fb88", size = 6722359, upload-time = "2024-05-29T17:56:07.991Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/61/53/e18c8c97d0b2724d85c9830477e3ebea3acf1dcdc6deb344d5d9c93a9946/sympy-1.12.1-py3-none-any.whl", hash = "sha256:9b2cbc7f1a640289430e13d2a56f02f867a1da0190f2f99d8968c2f74da0e515", size = 5743129, upload-time = "2024-05-29T17:56:05.445Z" }, +] + +[[package]] +name = "sympy" +version = "1.14.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.13.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'darwin')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'darwin')", + "(python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform == 'darwin')", + "(python_full_version == '3.10.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'darwin')", + "(python_full_version == '3.9.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform == 'darwin')", +] +dependencies = [ + { name = "mpmath", marker = "python_full_version >= '3.9'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, +] + +[[package]] +name = "threadpoolctl" +version = "3.5.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/bd/55/b5148dcbf72f5cde221f8bfe3b6a540da7aa1842f6b491ad979a6c8b84af/threadpoolctl-3.5.0.tar.gz", hash = "sha256:082433502dd922bf738de0d8bcc4fdcbf0979ff44c42bd40f5af8a282f6fa107", size = 41936, upload-time = "2024-04-29T13:50:16.544Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4b/2c/ffbf7a134b9ab11a67b0cf0726453cedd9c5043a4fe7a35d1cefa9a1bcfb/threadpoolctl-3.5.0-py3-none-any.whl", hash = "sha256:56c1e26c150397e58c4926da8eeee87533b1e32bef131bd4bf6a2f45f3185467", size = 18414, upload-time = "2024-04-29T13:50:14.014Z" }, +] + +[[package]] +name = "threadpoolctl" +version = "3.6.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.13.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'darwin')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'darwin')", + "(python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform == 'darwin')", + "(python_full_version == '3.10.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'darwin')", + "(python_full_version == '3.9.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform == 'darwin')", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b7/4d/08c89e34946fce2aec4fbb45c9016efd5f4d7f24af8e5d93296e935631d8/threadpoolctl-3.6.0.tar.gz", hash = "sha256:8ab8b4aa3491d812b623328249fab5302a68d2d71745c8a4c719a2fcaba9f44e", size = 21274, upload-time = "2025-03-13T13:49:23.031Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/32/d5/f9a850d79b0851d1d4ef6456097579a9005b31fea68726a4ae5f2d82ddd9/threadpoolctl-3.6.0-py3-none-any.whl", hash = "sha256:43a0b8fd5a2928500110039e43a5eed8480b918967083ea48dc3ab9f13c4a7fb", size = 18638, upload-time = "2025-03-13T13:49:21.846Z" }, +] + +[[package]] +name = "torch" +version = "2.5.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] +dependencies = [ + { name = "filelock", version = "3.16.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "fsspec", version = "2025.3.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "jinja2", marker = "python_full_version < '3.9'" }, + { name = "networkx", version = "3.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "nvidia-cublas-cu12", version = "12.4.5.8", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-cupti-cu12", version = "12.4.127", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvrtc-cu12", version = "12.4.127", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-runtime-cu12", version = "12.4.127", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cudnn-cu12", version = "9.1.0.70", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cufft-cu12", version = "11.2.1.3", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-curand-cu12", version = "10.3.5.147", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusolver-cu12", version = "11.6.1.9", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparse-cu12", version = "12.3.1.170", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu12", version = "2.21.5", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", version = "12.4.127", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvtx-cu12", version = "12.4.127", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "sympy", version = "1.12.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "triton", version = "3.1.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "typing-extensions", version = "4.13.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2a/ef/834af4a885b31a0b32fff2d80e1e40f771e1566ea8ded55347502440786a/torch-2.5.1-cp310-cp310-manylinux1_x86_64.whl", hash = "sha256:71328e1bbe39d213b8721678f9dcac30dfc452a46d586f1d514a6aa0a99d4744", size = 906446312, upload-time = "2024-10-29T17:33:38.045Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/69/f0/46e74e0d145f43fa506cb336eaefb2d240547e4ce1f496e442711093ab25/torch-2.5.1-cp310-cp310-manylinux2014_aarch64.whl", hash = "sha256:34bfa1a852e5714cbfa17f27c49d8ce35e1b7af5608c4bc6e81392c352dbc601", size = 91919522, upload-time = "2024-10-29T17:39:08.74Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a5/13/1eb674c8efbd04d71e4a157ceba991904f633e009a584dd65dccbafbb648/torch-2.5.1-cp310-cp310-win_amd64.whl", hash = "sha256:32a037bd98a241df6c93e4c789b683335da76a2ac142c0973675b715102dc5fa", size = 203088048, upload-time = "2024-10-29T17:34:10.913Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a9/9d/e0860474ee0ff8f6ef2c50ec8f71a250f38d78a9b9df9fd241ad3397a65b/torch-2.5.1-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:23d062bf70776a3d04dbe74db950db2a5245e1ba4f27208a87f0d743b0d06e86", size = 63877046, upload-time = "2024-10-29T17:34:19.174Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d1/35/e8b2daf02ce933e4518e6f5682c72fd0ed66c15910ea1fb4168f442b71c4/torch-2.5.1-cp311-cp311-manylinux1_x86_64.whl", hash = "sha256:de5b7d6740c4b636ef4db92be922f0edc425b65ed78c5076c43c42d362a45457", size = 906474467, upload-time = "2024-10-29T17:38:49.832Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/40/04/bd91593a4ca178ece93ca55f27e2783aa524aaccbfda66831d59a054c31e/torch-2.5.1-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:340ce0432cad0d37f5a31be666896e16788f1adf8ad7be481196b503dad675b9", size = 91919450, upload-time = "2024-10-29T17:37:26.693Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/0d/4a/e51420d46cfc90562e85af2fee912237c662ab31140ab179e49bd69401d6/torch-2.5.1-cp311-cp311-win_amd64.whl", hash = "sha256:603c52d2fe06433c18b747d25f5c333f9c1d58615620578c326d66f258686f9a", size = 203098237, upload-time = "2024-10-29T17:36:11.731Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d0/db/5d9cbfbc7968d79c5c09a0bc0bc3735da079f2fd07cc10498a62b320a480/torch-2.5.1-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:31f8c39660962f9ae4eeec995e3049b5492eb7360dd4f07377658ef4d728fa4c", size = 63884466, upload-time = "2024-10-29T17:33:02.899Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8b/5c/36c114d120bfe10f9323ed35061bc5878cc74f3f594003854b0ea298942f/torch-2.5.1-cp312-cp312-manylinux1_x86_64.whl", hash = "sha256:ed231a4b3a5952177fafb661213d690a72caaad97d5824dd4fc17ab9e15cec03", size = 906389343, upload-time = "2024-10-29T17:37:06.758Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6d/69/d8ada8b6e0a4257556d5b4ddeb4345ea8eeaaef3c98b60d1cca197c7ad8e/torch-2.5.1-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:3f4b7f10a247e0dcd7ea97dc2d3bfbfc90302ed36d7f3952b0008d0df264e697", size = 91811673, upload-time = "2024-10-29T17:32:42.789Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5f/ba/607d013b55b9fd805db2a5c2662ec7551f1910b4eef39653eeaba182c5b2/torch-2.5.1-cp312-cp312-win_amd64.whl", hash = "sha256:73e58e78f7d220917c5dbfad1a40e09df9929d3b95d25e57d9f8558f84c9a11c", size = 203046841, upload-time = "2024-10-29T17:35:48.665Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/57/6c/bf52ff061da33deb9f94f4121fde7ff3058812cb7d2036c97bc167793bd1/torch-2.5.1-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:8c712df61101964eb11910a846514011f0b6f5920c55dbf567bff8a34163d5b1", size = 63858109, upload-time = "2024-10-29T17:36:21.973Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/69/72/20cb30f3b39a9face296491a86adb6ff8f1a47a897e4d14667e6cf89d5c3/torch-2.5.1-cp313-cp313-manylinux1_x86_64.whl", hash = "sha256:9b61edf3b4f6e3b0e0adda8b3960266b9009d02b37555971f4d1c8f7a05afed7", size = 906393265, upload-time = "2024-10-29T17:35:06.866Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a9/18/81c399e8f4f1580d34bf99d827cb5fb5cf7a18a266bb5d30ca3ec2e89ba6/torch-2.5.1-cp39-cp39-manylinux1_x86_64.whl", hash = "sha256:1f3b7fb3cf7ab97fae52161423f81be8c6b8afac8d9760823fd623994581e1a3", size = 906479005, upload-time = "2024-10-29T17:39:51.513Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5d/86/1c4b168d52cddb8d17952a7b5b25f69ef0da1fc34de1223d73d0d9db1801/torch-2.5.1-cp39-cp39-manylinux2014_aarch64.whl", hash = "sha256:7974e3dce28b5a21fb554b73e1bc9072c25dde873fa00d54280861e7a009d7dc", size = 91846074, upload-time = "2024-10-29T17:33:57.857Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/76/49/4a0a8b19ce8f9bf32fcab4e863c7e2366f519f9826c84ca250567b11a014/torch-2.5.1-cp39-cp39-win_amd64.whl", hash = "sha256:46c817d3ea33696ad3b9df5e774dba2257e9a4cd3c4a3afbf92f6bb13ac5ce2d", size = 203000888, upload-time = "2024-10-29T17:32:53.884Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/25/07/3548a7cfcf69d0eccec2ee79ee3913f1cdaadb27b36946774db86729ee47/torch-2.5.1-cp39-none-macosx_11_0_arm64.whl", hash = "sha256:8046768b7f6d35b85d101b4b38cba8aa2f3cd51952bc4c06a49580f2ce682291", size = 63876023, upload-time = "2024-10-29T17:32:34.081Z" }, +] + +[[package]] +name = "torch" +version = "2.8.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.13.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'darwin')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'darwin')", + "(python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform == 'darwin')", + "(python_full_version == '3.10.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'darwin')", + "(python_full_version == '3.9.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform == 'darwin')", +] +dependencies = [ + { name = "filelock", version = "3.19.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9'" }, + { name = "fsspec", version = "2025.7.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9'" }, + { name = "jinja2", marker = "python_full_version >= '3.9'" }, + { name = "networkx", version = "3.2.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, + { name = "networkx", version = "3.4.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.10.*'" }, + { name = "networkx", version = "3.5", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "nvidia-cublas-cu12", version = "12.8.4.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-cupti-cu12", version = "12.8.90", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvrtc-cu12", version = "12.8.93", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-runtime-cu12", version = "12.8.90", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cudnn-cu12", version = "9.10.2.21", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cufft-cu12", version = "11.3.3.83", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cufile-cu12", marker = "python_full_version >= '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-curand-cu12", version = "10.3.9.90", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusolver-cu12", version = "11.7.3.90", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparse-cu12", version = "12.5.8.93", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparselt-cu12", marker = "python_full_version >= '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu12", version = "2.27.3", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", version = "12.8.93", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvtx-cu12", version = "12.8.90", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "setuptools", version = "80.9.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "sympy", version = "1.14.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9'" }, + { name = "triton", version = "3.4.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "typing-extensions", version = "4.15.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.9'" }, +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/63/28/110f7274254f1b8476c561dada127173f994afa2b1ffc044efb773c15650/torch-2.8.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:0be92c08b44009d4131d1ff7a8060d10bafdb7ddcb7359ef8d8c5169007ea905", size = 102052793, upload-time = "2025-08-06T14:53:15.852Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/70/1c/58da560016f81c339ae14ab16c98153d51c941544ae568da3cb5b1ceb572/torch-2.8.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:89aa9ee820bb39d4d72b794345cccef106b574508dd17dbec457949678c76011", size = 888025420, upload-time = "2025-08-06T14:54:18.014Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/70/87/f69752d0dd4ba8218c390f0438130c166fa264a33b7025adb5014b92192c/torch-2.8.0-cp310-cp310-win_amd64.whl", hash = "sha256:e8e5bf982e87e2b59d932769938b698858c64cc53753894be25629bdf5cf2f46", size = 241363614, upload-time = "2025-08-06T14:53:31.496Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ef/d6/e6d4c57e61c2b2175d3aafbfb779926a2cfd7c32eeda7c543925dceec923/torch-2.8.0-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:a3f16a58a9a800f589b26d47ee15aca3acf065546137fc2af039876135f4c760", size = 73611154, upload-time = "2025-08-06T14:53:10.919Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8f/c4/3e7a3887eba14e815e614db70b3b529112d1513d9dae6f4d43e373360b7f/torch-2.8.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:220a06fd7af8b653c35d359dfe1aaf32f65aa85befa342629f716acb134b9710", size = 102073391, upload-time = "2025-08-06T14:53:20.937Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5a/63/4fdc45a0304536e75a5e1b1bbfb1b56dd0e2743c48ee83ca729f7ce44162/torch-2.8.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:c12fa219f51a933d5f80eeb3a7a5d0cbe9168c0a14bbb4055f1979431660879b", size = 888063640, upload-time = "2025-08-06T14:55:05.325Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/84/57/2f64161769610cf6b1c5ed782bd8a780e18a3c9d48931319f2887fa9d0b1/torch-2.8.0-cp311-cp311-win_amd64.whl", hash = "sha256:8c7ef765e27551b2fbfc0f41bcf270e1292d9bf79f8e0724848b1682be6e80aa", size = 241366752, upload-time = "2025-08-06T14:53:38.692Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a4/5e/05a5c46085d9b97e928f3f037081d3d2b87fb4b4195030fc099aaec5effc/torch-2.8.0-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:5ae0524688fb6707c57a530c2325e13bb0090b745ba7b4a2cd6a3ce262572916", size = 73621174, upload-time = "2025-08-06T14:53:25.44Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/49/0c/2fd4df0d83a495bb5e54dca4474c4ec5f9c62db185421563deeb5dabf609/torch-2.8.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:e2fab4153768d433f8ed9279c8133a114a034a61e77a3a104dcdf54388838705", size = 101906089, upload-time = "2025-08-06T14:53:52.631Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/99/a8/6acf48d48838fb8fe480597d98a0668c2beb02ee4755cc136de92a0a956f/torch-2.8.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:b2aca0939fb7e4d842561febbd4ffda67a8e958ff725c1c27e244e85e982173c", size = 887913624, upload-time = "2025-08-06T14:56:44.33Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/af/8a/5c87f08e3abd825c7dfecef5a0f1d9aa5df5dd0e3fd1fa2f490a8e512402/torch-2.8.0-cp312-cp312-win_amd64.whl", hash = "sha256:2f4ac52f0130275d7517b03a33d2493bab3693c83dcfadf4f81688ea82147d2e", size = 241326087, upload-time = "2025-08-06T14:53:46.503Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/be/66/5c9a321b325aaecb92d4d1855421e3a055abd77903b7dab6575ca07796db/torch-2.8.0-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:619c2869db3ada2c0105487ba21b5008defcc472d23f8b80ed91ac4a380283b0", size = 73630478, upload-time = "2025-08-06T14:53:57.144Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/10/4e/469ced5a0603245d6a19a556e9053300033f9c5baccf43a3d25ba73e189e/torch-2.8.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:2b2f96814e0345f5a5aed9bf9734efa913678ed19caf6dc2cddb7930672d6128", size = 101936856, upload-time = "2025-08-06T14:54:01.526Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/16/82/3948e54c01b2109238357c6f86242e6ecbf0c63a1af46906772902f82057/torch-2.8.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:65616ca8ec6f43245e1f5f296603e33923f4c30f93d65e103d9e50c25b35150b", size = 887922844, upload-time = "2025-08-06T14:55:50.78Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e3/54/941ea0a860f2717d86a811adf0c2cd01b3983bdd460d0803053c4e0b8649/torch-2.8.0-cp313-cp313-win_amd64.whl", hash = "sha256:659df54119ae03e83a800addc125856effda88b016dfc54d9f65215c3975be16", size = 241330968, upload-time = "2025-08-06T14:54:45.293Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/de/69/8b7b13bba430f5e21d77708b616f767683629fc4f8037564a177d20f90ed/torch-2.8.0-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:1a62a1ec4b0498930e2543535cf70b1bef8c777713de7ceb84cd79115f553767", size = 73915128, upload-time = "2025-08-06T14:54:34.769Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/15/0e/8a800e093b7f7430dbaefa80075aee9158ec22e4c4fc3c1a66e4fb96cb4f/torch-2.8.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:83c13411a26fac3d101fe8035a6b0476ae606deb8688e904e796a3534c197def", size = 102020139, upload-time = "2025-08-06T14:54:39.047Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/4a/15/5e488ca0bc6162c86a33b58642bc577c84ded17c7b72d97e49b5833e2d73/torch-2.8.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:8f0a9d617a66509ded240add3754e462430a6c1fc5589f86c17b433dd808f97a", size = 887990692, upload-time = "2025-08-06T14:56:18.286Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b4/a8/6a04e4b54472fc5dba7ca2341ab219e529f3c07b6941059fbf18dccac31f/torch-2.8.0-cp313-cp313t-win_amd64.whl", hash = "sha256:a7242b86f42be98ac674b88a4988643b9bc6145437ec8f048fea23f72feb5eca", size = 241603453, upload-time = "2025-08-06T14:55:22.945Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/04/6e/650bb7f28f771af0cb791b02348db8b7f5f64f40f6829ee82aa6ce99aabe/torch-2.8.0-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:7b677e17f5a3e69fdef7eb3b9da72622f8d322692930297e4ccb52fefc6c8211", size = 73632395, upload-time = "2025-08-06T14:55:28.645Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5b/b0/a321f27270049baa12f5c3fb0d6ceea005634787e3af9a8d75dce8306b0a/torch-2.8.0-cp39-cp39-manylinux_2_28_aarch64.whl", hash = "sha256:da6afa31c13b669d4ba49d8a2169f0db2c3ec6bec4af898aa714f401d4c38904", size = 102059214, upload-time = "2025-08-06T14:55:33.433Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/fd/dd/1630cb51b10d3d2e97db95e5a84c32def81fc26b005bce6fc880b0e6db81/torch-2.8.0-cp39-cp39-manylinux_2_28_x86_64.whl", hash = "sha256:06fcee8000e5c62a9f3e52a688b9c5abb7c6228d0e56e3452983416025c41381", size = 888024302, upload-time = "2025-08-06T14:57:28.23Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/b9/dc/1f1f621afe15e3c496e1e8f94f8903f75f87e7d642d5a985e92210cc208d/torch-2.8.0-cp39-cp39-win_amd64.whl", hash = "sha256:5128fe752a355d9308e56af1ad28b15266fe2da5948660fad44de9e3a9e36e8c", size = 241249338, upload-time = "2025-08-06T14:57:05.669Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ae/95/ae26263aceb3d57b821179f827d0e321373ed49423e603dd5906ab14a730/torch-2.8.0-cp39-none-macosx_11_0_arm64.whl", hash = "sha256:e9f071f5b52a9f6970dc8a919694b27a91ae9dc08898b2b988abbef5eddfd1ae", size = 73610795, upload-time = "2025-08-06T14:57:11.513Z" }, +] + +[[package]] +name = "torch-complex" +version = "0.4.4" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +dependencies = [ + { name = "numpy", version = "1.24.4", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "numpy", version = "2.0.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.10'" }, + { name = "packaging" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/bf/2b/17cb15a383cf2135330371e034d13b9043dc6d8bd07c871b5aa3064fbed1/torch_complex-0.4.4.tar.gz", hash = "sha256:4153fd6b24a0bad689e6f193bfbd00f38283b1890d808bef684ddc6d1f63fd3f", size = 10025, upload-time = "2024-06-28T07:10:28.136Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f4/c5/9b4d756a7ada951e9b17dcc636f98ed1073c737ae809b150ef408afb6298/torch_complex-0.4.4-py3-none-any.whl", hash = "sha256:6ab4ecd4f3a16e3adb70a7f7cd2e769a9dfd07d7a8e27d04ff9c621ebbe34b13", size = 9125, upload-time = "2024-06-28T07:10:26.651Z" }, +] + +[[package]] +name = "torchaudio" +version = "2.5.1+cpu" +source = { registry = "https://download.pytorch.org/whl/cpu" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] +dependencies = [ + { name = "torch", version = "2.5.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, +] +wheels = [ + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.5.1%2Bcpu-cp310-cp310-linux_x86_64.whl", hash = "sha256:c64af43548713e5abc3e9e3b5f33a2a47c57122ee953e0b0cb102f7855f8d017" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.5.1%2Bcpu-cp310-cp310-win_amd64.whl", hash = "sha256:8904a77ad311992d3dd3dae6556ea023fbb291b10b8698fecd1dec10bc189d6c" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.5.1%2Bcpu-cp311-cp311-linux_x86_64.whl", hash = "sha256:209f8e7efff818d2066f49c27c8097f2246f8bf6fd743c69bcf8dce80c63ee80" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.5.1%2Bcpu-cp311-cp311-win_amd64.whl", hash = "sha256:e08047b9f1997bd303ea6528e5b8009719665a55516a49b6acc359dac20c474d" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.5.1%2Bcpu-cp312-cp312-linux_x86_64.whl", hash = "sha256:b7fbd9c264dcbe28efb061364c76d3770eb13ae692d2982949b583edfb9ed7f5" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.5.1%2Bcpu-cp312-cp312-win_amd64.whl", hash = "sha256:68b58c1dae7bb738df8a36f0d147d2c9c69e56a35ff40dedb84ac30aff2e59e1" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.5.1%2Bcpu-cp39-cp39-linux_x86_64.whl", hash = "sha256:63518869ddb6bbf37a305bb0c085552c230e130ea4640fd71db17b6689586aa5" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.5.1%2Bcpu-cp39-cp39-win_amd64.whl", hash = "sha256:ea2c7fd25311ea47fd78abf0ee6ffbc4820d785773d0bc5ac87f82912851643f" }, +] + +[[package]] +name = "torchaudio" +version = "2.8.0" +source = { registry = "https://download.pytorch.org/whl/cpu" } +resolution-markers = [ + "(python_full_version == '3.13.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'darwin')", + "(python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'darwin')", + "(python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform == 'darwin')", + "(python_full_version == '3.10.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'darwin')", + "(python_full_version == '3.9.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform == 'darwin')", +] +dependencies = [ + { name = "torch", version = "2.8.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "(python_full_version >= '3.9' and python_full_version < '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.9' and sys_platform == 'darwin')" }, +] +wheels = [ + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:c2f44cf279f673cfcdd8f576c349eee8bedf8caab351a5dd78b32970cc34a212" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:d3c1b85b26a09832d139f6d6da6b66caeb51d2e16e08f8587665c44a9e1aa8f9" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:c9276857d241c6de257af765c0f51fc011af38cb725401495121b280913007cf" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:4573c6042950c20278e3608a9a38050ba0bc72e0049e1bbfd249caf859a8029b" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ddef94bf181e6447cbb05f38beaca8f6c5bb8d2b9ddced1aa3452025b9fc70d3" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:862e2e40bf09d865e5df080a84c1a39bbcef40e43140f4b1737eb3a389d3b38f" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:f851d32e94ca05e470f0c60e25726ec1e0eb71cb2ca5a0206b7fd03272ccc3c8" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:09535a9b727c0793cd07c1ace99f3f353626281bcc3e30c2f2314e3ebc9d3f96" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:68df9c9068984edff8065c2b6656725e6114fe89281b0cf122c7505305fc98a4" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:1951f10ed092f2dda57634f6a3950ef21c9d9352551aa84a9fccd51bbda18095" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:522289e2cd57e79401fd5ccae9b1bc0ff2e47f3529092adf5acf57427ea0c6a9" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0-cp39-cp39-manylinux_2_28_aarch64.whl", hash = "sha256:ef7f7ffa828b8d8ba5d3a569b825fc0469688e1e8962bf6577d538bd8af1387d" }, +] + +[[package]] +name = "torchaudio" +version = "2.8.0+cpu" +source = { registry = "https://download.pytorch.org/whl/cpu" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", +] +dependencies = [ + { name = "torch", version = "2.8.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.9' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.9' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.9' and sys_platform != 'darwin' and sys_platform != 'linux')" }, +] +wheels = [ + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0%2Bcpu-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:c955835e470ebbde03d7d54ca5d8ba5722138bbfd66cfb86845234b3a5b9f9fa" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0%2Bcpu-cp310-cp310-win_amd64.whl", hash = "sha256:e9e68f16f1afe108f0cb1c7d636d0242fdc43cbbcaab222a72a373b9d2799134" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0%2Bcpu-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:e54bd7fc9472019308097d99102df9acee22aa2451ae808d27840bc874320292" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0%2Bcpu-cp311-cp311-win_amd64.whl", hash = "sha256:db37df7eee906f8fe0a639fdc673f3541cb2e173169b16d4133447eb922d1938" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0%2Bcpu-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:9377faee65a290578280ac7f4884c3586253dac2ca28c60f458ff6efe86a6b05" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0%2Bcpu-cp312-cp312-win_amd64.whl", hash = "sha256:9b302192b570657c1cc787a4d487ae4bbb7f2aab1c01b1fcc46757e7f86f391e" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0%2Bcpu-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:e1b1f530e8b71b1d079e23db45a0e621709061710ef8540aae8280aa039554ee" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0%2Bcpu-cp313-cp313-win_amd64.whl", hash = "sha256:0c2d081e24204768e636cbf05e1377c8a6964b8ed6fa3aa5092ba9af9bbc19c5" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0%2Bcpu-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:89c2d04fe1cb7c31eb042f7b36e1ce8e2afacf769ecd5f216527e184e4857099" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0%2Bcpu-cp313-cp313t-win_amd64.whl", hash = "sha256:ab4653da31dc37f0a643f41f4da8bee647a8686bacf12d3929cac8aead186811" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0%2Bcpu-cp39-cp39-manylinux_2_28_x86_64.whl", hash = "sha256:e856b1abb280e1d961bdc12476bd38fc7eab8af720f9c903c95998dda069ae4c" }, + { url = "https://download.pytorch.org/whl/cpu/torchaudio-2.8.0%2Bcpu-cp39-cp39-win_amd64.whl", hash = "sha256:7e9a06f6dc73f98aff1a5540f8d6103b66e4c945c1d94612087954905f221171" }, +] + +[[package]] +name = "tpu-perf" +version = "1.2.35" +source = { path = "bmodel/tpu_perf-1.2.35-py3-none-manylinux2014_aarch64.whl" } +dependencies = [ + { name = "lmdb" }, + { name = "numpy", version = "1.24.4", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "numpy", version = "2.0.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.10'" }, + { name = "openpyxl" }, + { name = "panda" }, + { name = "protobuf" }, + { name = "psutil" }, + { name = "pyyaml" }, +] +wheels = [ + { filename = "tpu_perf-1.2.35-py3-none-manylinux2014_aarch64.whl", hash = "sha256:0608ae719c1701a022d6881c9d86c2cf290ce0de0ffeeff40fb089088a9e50c9" }, +] + +[package.metadata] +requires-dist = [ + { name = "lmdb" }, + { name = "numpy" }, + { name = "openpyxl" }, + { name = "panda" }, + { name = "protobuf", specifier = "<3.21" }, + { name = "psutil" }, + { name = "pyyaml" }, +] + +[[package]] +name = "tqdm" +version = "4.67.1" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a8/4b/29b4ef32e036bb34e4ab51796dd745cdba7ed47ad142a9f4a1eb8e0c744d/tqdm-4.67.1.tar.gz", hash = "sha256:f8aef9c52c08c13a65f30ea34f4e5aac3fd1a34959879d7e59e63027286627f2", size = 169737, upload-time = "2024-11-24T20:12:22.481Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d0/30/dc54f88dd4a2b5dc8a0279bdd7270e735851848b762aeb1c1184ed1f6b14/tqdm-4.67.1-py3-none-any.whl", hash = "sha256:26445eca388f82e72884e0d580d5464cd801a3ea01e63e5601bdff9ba6a48de2", size = 78540, upload-time = "2024-11-24T20:12:19.698Z" }, +] + +[[package]] +name = "triton" +version = "3.1.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", +] +dependencies = [ + { name = "filelock", version = "3.16.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9' and sys_platform != 'darwin'" }, +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/98/29/69aa56dc0b2eb2602b553881e34243475ea2afd9699be042316842788ff5/triton-3.1.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6b0dd10a925263abbe9fa37dcde67a5e9b2383fc269fdf59f5657cac38c5d1d8", size = 209460013, upload-time = "2024-10-14T16:05:32.106Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/86/17/d9a5cf4fcf46291856d1e90762e36cbabd2a56c7265da0d1d9508c8e3943/triton-3.1.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0f34f6e7885d1bf0eaaf7ba875a5f0ce6f3c13ba98f9503651c1e6dc6757ed5c", size = 209506424, upload-time = "2024-10-14T16:05:42.337Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/78/eb/65f5ba83c2a123f6498a3097746607e5b2f16add29e36765305e4ac7fdd8/triton-3.1.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c8182f42fd8080a7d39d666814fa36c5e30cc00ea7eeeb1a2983dbb4c99a0fdc", size = 209551444, upload-time = "2024-10-14T16:05:53.433Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/15/3c/e972ac0dd0f35ba5fb7058152dd52127a225f579eba2d7527eb1ffb3891a/triton-3.1.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6dadaca7fc24de34e180271b5cf864c16755702e9f63a16f62df714a8099126a", size = 209434868, upload-time = "2024-10-14T16:06:03.249Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c4/69/57e0fed438d547524e08bfedc587078314176ad1c15c8be904d3f03149ec/triton-3.1.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:aafa9a20cd0d9fee523cd4504aa7131807a864cd77dcf6efe7e981f18b8c6c11", size = 209460480, upload-time = "2024-10-14T16:06:13.116Z" }, +] + +[[package]] +name = "triton" +version = "3.4.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", +] +dependencies = [ + { name = "importlib-metadata", version = "8.7.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')" }, + { name = "setuptools", version = "80.9.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.9' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.9' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.9' and sys_platform != 'darwin' and sys_platform != 'linux')" }, +] +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/62/ee/0ee5f64a87eeda19bbad9bc54ae5ca5b98186ed00055281fd40fb4beb10e/triton-3.4.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7ff2785de9bc02f500e085420273bb5cc9c9bb767584a4aa28d6e360cec70128", size = 155430069, upload-time = "2025-07-30T19:58:21.715Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7d/39/43325b3b651d50187e591eefa22e236b2981afcebaefd4f2fc0ea99df191/triton-3.4.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7b70f5e6a41e52e48cfc087436c8a28c17ff98db369447bcaff3b887a3ab4467", size = 155531138, upload-time = "2025-07-30T19:58:29.908Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/d0/66/b1eb52839f563623d185f0927eb3530ee4d5ffe9d377cdaf5346b306689e/triton-3.4.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:31c1d84a5c0ec2c0f8e8a072d7fd150cab84a9c239eaddc6706c081bfae4eb04", size = 155560068, upload-time = "2025-07-30T19:58:37.081Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/30/7b/0a685684ed5322d2af0bddefed7906674f67974aa88b0fae6e82e3b766f6/triton-3.4.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:00be2964616f4c619193cb0d1b29a99bd4b001d7dc333816073f92cf2a8ccdeb", size = 155569223, upload-time = "2025-07-30T19:58:44.017Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/20/63/8cb444ad5cdb25d999b7d647abac25af0ee37d292afc009940c05b82dda0/triton-3.4.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7936b18a3499ed62059414d7df563e6c163c5e16c3773678a3ee3d417865035d", size = 155659780, upload-time = "2025-07-30T19:58:51.171Z" }, + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/12/34/1251beb5a3cb93f3950ebe68732752014646003ef6eb11eb5f1a37ca78cd/triton-3.4.0-cp39-cp39-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:98e5c1442eaeabae2e2452ae765801bd53cd4ce873cab0d1bdd59a32ab2d9397", size = 155430799, upload-time = "2025-07-30T19:58:57.664Z" }, +] + +[[package]] +name = "typing-extensions" +version = "4.13.2" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/f6/37/23083fcd6e35492953e8d2aaaa68b860eb422b34627b13f2ce3eb6106061/typing_extensions-4.13.2.tar.gz", hash = "sha256:e6c81219bd689f51865d9e372991c540bda33a0379d5573cddb9a3a23f7caaef", size = 106967, upload-time = "2025-04-10T14:19:05.416Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/8b/54/b1ae86c0973cc6f0210b53d508ca3641fb6d0c56823f288d108bc7ab3cc8/typing_extensions-4.13.2-py3-none-any.whl", hash = "sha256:a439e7c04b49fec3e5d3e2beaa21755cadbbdc391694e28ccdd36ca4a1408f8c", size = 45806, upload-time = "2025-04-10T14:19:03.967Z" }, +] + +[[package]] +name = "typing-extensions" +version = "4.15.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.13.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'darwin')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'darwin')", + "(python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform == 'darwin')", + "(python_full_version == '3.10.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'darwin')", + "(python_full_version == '3.9.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform == 'darwin')", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/72/94/1a15dd82efb362ac84269196e94cf00f187f7ed21c242792a923cdb1c61f/typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466", size = 109391, upload-time = "2025-08-25T13:49:26.313Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" }, +] + +[[package]] +name = "umap-learn" +version = "0.5.7" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] +dependencies = [ + { name = "numba", version = "0.58.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "numpy", version = "1.24.4", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "pynndescent", marker = "python_full_version < '3.9'" }, + { name = "scikit-learn", version = "1.3.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "scipy", version = "1.10.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version < '3.9'" }, + { name = "tqdm", marker = "python_full_version < '3.9'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6f/d4/9ed627905f7993349671283b3c5bf2d9f543ef79229fa1c7e01324eb900c/umap-learn-0.5.7.tar.gz", hash = "sha256:b2a97973e4c6ffcebf241100a8de589a4c84126a832ab40f296c6d9fcc5eb19e", size = 92680, upload-time = "2024-10-28T18:05:57.093Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/3c/8f/671c0e1f2572ba625cbcc1faeba9435e00330c3d6962858711445cf1e817/umap_learn-0.5.7-py3-none-any.whl", hash = "sha256:6a7e0be2facfa365a5ed6588447102bdbef32a0ef449535c25c97ea7e680073c", size = 88815, upload-time = "2024-10-28T18:05:55.333Z" }, +] + +[[package]] +name = "umap-learn" +version = "0.5.9.post2" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.13.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'darwin')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'darwin')", + "(python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform == 'darwin')", + "(python_full_version == '3.10.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'darwin')", + "(python_full_version == '3.9.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform == 'darwin')", +] +dependencies = [ + { name = "numba", version = "0.60.0", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, + { name = "numba", version = "0.61.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.10'" }, + { name = "numpy", version = "2.0.2", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.10'" }, + { name = "pynndescent", marker = "python_full_version >= '3.9'" }, + { name = "scikit-learn", version = "1.6.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, + { name = "scikit-learn", version = "1.7.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.10'" }, + { name = "scipy", version = "1.13.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.9.*'" }, + { name = "scipy", version = "1.15.3", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version == '3.10.*'" }, + { name = "scipy", version = "1.16.1", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "tqdm", marker = "python_full_version >= '3.9'" }, +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/5f/ee/6bc65bd375c812026a7af63fe9d09d409382120aff25f2152f1ba12af5ec/umap_learn-0.5.9.post2.tar.gz", hash = "sha256:bdf60462d779bd074ce177a0714ced17e6d161285590fa487f3f9548dd3c31c9", size = 95441, upload-time = "2025-07-03T00:18:02.479Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/6b/b1/c24deeda9baf1fd491aaad941ed89e0fed6c583a117fd7b79e0a33a1e6c0/umap_learn-0.5.9.post2-py3-none-any.whl", hash = "sha256:fbe51166561e0e7fab00ef3d516ac2621243b8d15cf4bef9f656d701736b16a0", size = 90146, upload-time = "2025-07-03T00:18:01.042Z" }, +] + +[[package]] +name = "urllib3" +version = "2.2.3" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ed/63/22ba4ebfe7430b76388e7cd448d5478814d3032121827c12a2cc287e2260/urllib3-2.2.3.tar.gz", hash = "sha256:e7d814a81dad81e6caf2ec9fdedb284ecc9c73076b62654547cc64ccdcae26e9", size = 300677, upload-time = "2024-09-12T10:52:18.401Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ce/d9/5f4c13cecde62396b0d3fe530a50ccea91e7dfc1ccf0e09c228841bb5ba8/urllib3-2.2.3-py3-none-any.whl", hash = "sha256:ca899ca043dcb1bafa3e262d73aa25c465bfb49e0bd9dd5d59f1d0acba2f8fac", size = 126338, upload-time = "2024-09-12T10:52:16.589Z" }, +] + +[[package]] +name = "urllib3" +version = "2.5.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version >= '3.14' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.13' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.13.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'darwin')", + "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.10.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.10.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", + "(python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'darwin')", + "(python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform == 'darwin')", + "(python_full_version == '3.10.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.10.*' and sys_platform == 'darwin')", + "(python_full_version == '3.9.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform == 'darwin')", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/15/22/9ee70a2574a4f4599c47dd506532914ce044817c7752a79b6a51286319bc/urllib3-2.5.0.tar.gz", hash = "sha256:3fc47733c7e419d4bc3f6b3dc2b4f890bb743906a30d56ba4a5bfa4bbff92760", size = 393185, upload-time = "2025-06-18T14:07:41.644Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/a7/c2/fe1e52489ae3122415c51f387e221dd0773709bad6c6cdaa599e8a2c5185/urllib3-2.5.0-py3-none-any.whl", hash = "sha256:e6b01673c0fa6a13e374b50871808eb3bf7046c4b125b216f6bf1cc604cff0dc", size = 129795, upload-time = "2025-06-18T14:07:40.39Z" }, +] + +[[package]] +name = "zipp" +version = "3.20.2" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "python_full_version < '3.9' and sys_platform != 'darwin'", + "python_full_version < '3.9' and sys_platform == 'darwin'", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/54/bf/5c0000c44ebc80123ecbdddba1f5dcd94a5ada602a9c225d84b5aaa55e86/zipp-3.20.2.tar.gz", hash = "sha256:bc9eb26f4506fda01b81bcde0ca78103b6e62f991b381fec825435c836edbc29", size = 24199, upload-time = "2024-09-13T13:44:16.101Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/62/8b/5ba542fa83c90e09eac972fc9baca7a88e7e7ca4b221a89251954019308b/zipp-3.20.2-py3-none-any.whl", hash = "sha256:a817ac80d6cf4b23bf7f2828b7cabf326f15a001bea8b1f9b49631780ba28350", size = 9200, upload-time = "2024-09-13T13:44:14.38Z" }, +] + +[[package]] +name = "zipp" +version = "3.23.0" +source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" } +resolution-markers = [ + "(python_full_version == '3.9.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.9.*' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (python_full_version == '3.9.*' and sys_platform != 'darwin' and sys_platform != 'linux')", +] +sdist = { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/e3/02/0f2892c661036d50ede074e376733dca2ae7c6eb617489437771209d4180/zipp-3.23.0.tar.gz", hash = "sha256:a07157588a12518c9d4034df3fbbee09c814741a33ff63c05fa29d26a2404166", size = 25547, upload-time = "2025-06-08T17:06:39.4Z" } +wheels = [ + { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/2e/54/647ade08bf0db230bfea292f893923872fd20be6ac6f53b2b936ba839d75/zipp-3.23.0-py3-none-any.whl", hash = "sha256:071652d6115ed432f5ce1d34c336c0adfd6a884660d1e9712a256d3d3bd4b14e", size = 10276, upload-time = "2025-06-08T17:06:38.034Z" }, +]