Upload folder using huggingface_hub

MyEncoderASR.py

@@ -1,7 +1,10 @@
 from speechbrain.inference.ASR import EncoderASR
+from speechbrain.decoders.ctc import TorchAudioCTCPrefixBeamSearcher
+from speechbrain.decoders.ctc import CTCHypothesis
 import torch
 import speechbrain
 import functools
+import matplotlib.pyplot as plt
 
 class MyEncoderASR(EncoderASR):
     def transcribe_batch(self, wavs, wav_lens):
@@ -33,6 +36,7 @@ class MyEncoderASR(EncoderASR):
         with torch.no_grad():
             wav_lens = wav_lens.to(self.device)
             encoder_out = self.encode_batch(wavs, wav_lens)
+            # frame level logits.
             predictions = self.decoding_function(encoder_out, wav_lens)
             is_ctc_text_encoder_tokenizer = isinstance(
                 self.tokenizer, speechbrain.dataio.encoder.CTCTextEncoder
@@ -50,5 +54,156 @@ class MyEncoderASR(EncoderASR):
                     ]
             else:
                 predicted_words = [hyp[0].text for hyp in predictions]
-
         return predicted_words, predictions
+
+class MyCTCPrefixBeamSearcher(TorchAudioCTCPrefixBeamSearcher):
+    def decode_beams(self, log_probs, wav_len):
+        """Decode log_probs using TorchAudio CTC decoder.
+
+        If `using_cpu_decoder=True` then log_probs and wav_len are moved to CPU before decoding.
+        When using CUDA CTC decoder, the timestep information is not available. Therefore, the timesteps
+        in the returned hypotheses are set to None.
+
+        Make sure that the input are in the log domain. The decoder will fail to decode
+        logits or probabilities. The input should be the log probabilities of the CTC output.
+
+        Arguments
+        ---------
+        log_probs : torch.Tensor
+            The log probabilities of the input audio.
+            Shape: (batch_size, seq_length, vocab_size)
+        wav_len : torch.Tensor, default: None
+            The speechbrain-style relative length. Shape: (batch_size,)
+            If None, then the length of each audio is assumed to be seq_length.
+
+        Returns
+        -------
+        list of list of CTCHypothesis
+            The decoded hypotheses. The outer list is over the batch dimension, and the inner list is over the topk dimension.
+        """
+        if wav_len is not None:
+            wav_len = log_probs.size(1) * wav_len
+        else:
+            wav_len = torch.tensor(
+                [log_probs.size(1)] * log_probs.size(0),
+                device=log_probs.device,
+                dtype=torch.int32,
+            )
+
+        if wav_len.dtype != torch.int32:
+            wav_len = wav_len.to(torch.int32)
+
+        if log_probs.dtype != torch.float32:
+            raise ValueError("log_probs must be float32.")
+
+        # When using CPU decoder, we need to move the log_probs and wav_len to CPU
+        if self.using_cpu_decoder and log_probs.is_cuda:
+            log_probs = log_probs.cpu()
+
+        if self.using_cpu_decoder and wav_len.is_cuda:
+            wav_len = wav_len.cpu()
+
+        if not log_probs.is_contiguous():
+            raise RuntimeError("log_probs must be contiguous.")
+
+        results = self._ctc_decoder(log_probs, wav_len)
+
+        tokens_preds = []
+        words_preds = []
+        scores_preds = []
+        timesteps_preds = []
+
+        # over batch dim
+        for i in range(len(results)):
+            if self.using_cpu_decoder:
+                preds = [
+                    results[i][j].tokens.tolist()
+                    for j in range(len(results[i]))
+                ]
+                preds = [
+                    [self.tokens[token] for token in tokens] for tokens in preds
+                ]
+                tokens_preds.append(preds)
+
+                timesteps = [
+                    results[i][j].timesteps.tolist()
+                    for j in range(len(results[i]))
+                ]
+                timesteps_preds.append(timesteps)
+
+            else:
+                # no timesteps is available for CUDA CTC decoder
+                timesteps = [None for _ in range(len(results[i]))]
+                timesteps_preds.append(timesteps)
+
+                preds = [results[i][j].tokens for j in range(len(results[i]))]
+                preds = [
+                    [self.tokens[token] for token in tokens] for tokens in preds
+                ]
+                tokens_preds.append(preds)
+
+            words = [results[i][j].words for j in range(len(results[i]))]
+            words_preds.append(words)
+
+            scores = [results[i][j].score for j in range(len(results[i]))]
+            scores_preds.append(scores)
+
+        hyps = []
+        for (
+            batch_index,
+            (batch_text, batch_score, batch_timesteps),
+        ) in enumerate(zip(tokens_preds, scores_preds, timesteps_preds)):
+            hyps.append([])
+            for text, score, timestep in zip(
+                batch_text, batch_score, batch_timesteps
+            ):
+                hyps[batch_index].append(
+                    CTCHypothesis(
+                        text=text,
+                        last_lm_state=None,
+                        score=score,
+                        lm_score=score,
+                        text_frames=timestep,
+                    )
+                )
+        return hyps
+
+def plot_alignments(waveform, emission, tokens, timesteps, sample_rate):
+    t = torch.arange(waveform.size(0)) / sample_rate
+    ratio = waveform.size(0) / emission.size(1) / sample_rate
+
+    chars = []
+    words = []
+    word_start = None
+    for token, timestep in zip(tokens, timesteps * ratio):
+        if token == "|":
+            if word_start is not None:
+                words.append((word_start, timestep))
+            word_start = None
+        else:
+            chars.append((token, timestep))
+            if word_start is None:
+                word_start = timestep
+
+    num_axes = len(waveform) // sample_rate + 1
+    plt.figure(figsize=[num_axes*10, 5])
+    fig, axes = plt.subplots(num_axes, 1)
+
+    def _plot(ax, xlim):
+        ax.plot(t, waveform)
+        for token, timestep in chars:
+            ax.annotate(token.upper(), (timestep, 0.5))
+        for word_start, word_end in words:
+            ax.axvspan(word_start, word_end, alpha=0.1, color="red")
+        ax.set_ylim(-0.6, 0.7)
+        ax.set_yticks([0])
+        ax.grid(True, axis="y")
+        ax.set_xlim(xlim)
+
+    for i in range(0, num_axes):
+        _plot(axes[i], (i, i+1))
+    
+    axes[num_axes-1].set_xlabel("time (sec)")
+    fig.tight_layout()
+    
+    return fig

inference.yaml

@@ -67,8 +67,8 @@ training_target: "target" # "target": deduplicated canonical phoneme sequence; "
                           # "canonical"
                           # "perceived": deduplicated perceived phoneme sequence
 
-perceived_ssl: !apply:trainer.AutoSSLoader.AutoSSLLoader
-    model_name: !ref <perceived_ssl_model>
+perceived_ssl: !new:speechbrain.lobes.models.huggingface_transformers.wavlm.WavLM
+    source: "microsoft/wavlm-large"
     freeze: !ref <freeze_perceived_ssl>
     freeze_feature_extractor: !ref <freeze_perceived_feature_extractor>
     save_path: !ref <pretrained_models_path>

inference_verbose.yaml

@@ -0,0 +1,202 @@
+# Hyperparameters toggles
+prefix: ""
+
+## SSL features Selection
+pretrained_models_path: pretrained_models/
+# pretrained_models:
+# {
+#     "wav2vec2_base": "facebook/wav2vec2-base", # 768
+#     "hubert_base": "facebook/hubert-base-ls960", # 768
+#     "wavlm_base": "microsoft/wavlm-base", # 768
+#     "wavlm_base_plus": "microsoft/wavlm-base-plus", # 768
+#     "hubert_multilingual": "utter-project/mHuBERT-147", # 768
+#     "clap" : "laion/clap-htsat-fused", # 768
+#     "data2vec_base": "facebook/data2vec-audio-base", # 768
+    
+#     "wav2vec2_large": "facebook/wav2vec2-large", # 1024
+#     "hubert_large": "facebook/hubert-large-ls960", # 1024
+#     "wavlm_large": "microsoft/wavlm-large-plus", # 1024
+#     "data2vec_large": "facebook/data2vec-audio-large", #1024
+#     "whisper_medium": "openai/whisper-medium", # 1024
+    
+#     "whisper_large_v3_turbo": "openai/whisper-large-v3-turbo", # 1280
+# }
+
+
+# select pretrained SSL models
+perceived_ssl_model: "wavlm_large" # in pretrained_models
+canonical_ssl_model: Null
+
+# # models hidden size, varies by model
+ENCODER_DIM: 1024
+
+# # How to fuse the features
+feature_fusion: "mono"        # Options: "mono" for single ssl, "dual_ssl_enc" for dual ssl encoder, "dual_loss" for single SSL dual ssl loss
+blend_alpha: 0.5              # If using "blend" fusion
+
+# Input files
+# Data files
+data_folder_save: "./data"
+train_annotation: !ref <data_folder_save>/train-train.json
+valid_annotation: !ref <data_folder_save>/train-dev.json
+test_annotation: !ref <data_folder_save>/test.json
+# Extra data
+train_annotation_extra: !ref <data_folder_save>/train-train_with_extra.json
+use_extra_train_data: False
+
+evaluate_key: "PER" # use "mpd_f1_seq" for Transformer decoder path best mpd f1
+                            # "PER_seq" for Transformer decoder's best error rate
+                            # "PER" for ctc path best error rate
+                            # "mpd_f1" for ctc path best mpd f1
+max_save_models: 3 # Maximum number of saved models for each metrics 
+# generate training id for output folder
+# generate_training_id: !apply:trainer.generate_training_id.generate_training_id [!ref <perceived_ssl_model_id>, !ref <canonical_ssl_model_id>, !ref <feature_fusion>, !ref <prefix>]
+
+# output files
+output_folder: !ref exp_l2arctic/<perceived_ssl_model>_<canonical_ssl_model>_<feature_fusion>_<prefix>
+per_file: !ref <output_folder>/per.txt
+mpd_file: !ref <output_folder>/mpd.txt
+save_folder: !ref <output_folder>/save
+train_log: !ref <output_folder>/train_log.txt
+
+on_training_test_wer_folder: !ref <output_folder>/on_training_test_wer
+on_training_test_mpd_folder: !ref <output_folder>/on_training_test_mpd
+
+# Training Target 
+training_target: "target" # "target": deduplicated canonical phoneme sequence; "target_with_repeats": with repeats
+                          # "canonical"
+                          # "perceived": deduplicated perceived phoneme sequence
+
+# perceived_ssl: !apply:trainer.AutoSSLoader.AutoSSLLoader
+perceived_ssl: !new:speechbrain.lobes.models.huggingface_transformers.wavlm.WavLM
+    source: "microsoft/wavlm-large"
+    freeze: !ref <freeze_perceived_ssl>
+    freeze_feature_extractor: !ref <freeze_perceived_feature_extractor>
+    save_path: !ref <pretrained_models_path>
+    output_all_hiddens: False
+preceived_ssl_emb_layer: -1
+
+enc: !new:torch.nn.Sequential
+  - !new:speechbrain.lobes.models.VanillaNN.VanillaNN
+     input_shape: [null, null, !ref <ENCODER_DIM>]
+     activation: !ref <activation>
+     dnn_blocks: !ref <dnn_layers>
+     dnn_neurons: !ref <dnn_neurons>
+  - !new:torch.nn.LayerNorm
+     normalized_shape: !ref <dnn_neurons>
+
+ctc_lin:  !new:speechbrain.nnet.linear.Linear
+    input_size: !ref <dnn_neurons>
+    n_neurons: !ref <output_neurons>  # 40 phonemes + 1 blank + 1 err
+
+# Model parameters
+activation: !name:torch.nn.LeakyReLU
+dnn_layers: 2
+dnn_neurons: 384  
+freeze_perceived_ssl: False
+freeze_canonical_ssl: False
+freeze_perceived_feature_extractor: True  # freeze the CNN extractor in wav2vec
+freeze_canonical_feature_extractor: True         # Freeze Whisper encoder?
+
+log_softmax: !new:speechbrain.nnet.activations.Softmax
+    apply_log: True
+
+ctc_cost: !name:speechbrain.nnet.losses.ctc_loss
+    blank_index: !ref <blank_index>
+
+ctc_cost_mispro: !name:speechbrain.nnet.losses.ctc_loss
+    blank_index: !ref <blank_index>
+
+# Outputs
+output_neurons: 44 # l2arctic: 40phns(sil)+err+blank + eos + bos =44
+blank_index: 0
+
+model: !new:torch.nn.ModuleList
+    - [!ref <enc>, !ref <ctc_lin>]
+
+adam_opt_class: !name:torch.optim.Adam
+    lr: !ref <lr>
+
+pretrained_opt_class: !name:torch.optim.Adam
+    lr: !ref <lr_pretrained>
+
+checkpointer: !new:speechbrain.utils.checkpoints.Checkpointer
+    checkpoints_dir: !ref <save_folder>
+    recoverables:
+        model: !ref <model>
+        perceived_ssl: !ref <perceived_ssl>
+        counter: !ref <epoch_counter>
+    allow_partial_load: True
+# canonical_ssl: !ref <canonical_ssl>
+
+augmentation: !new:speechbrain.augment.time_domain.SpeedPerturb
+    orig_freq: !ref <sample_rate>
+    speeds: [95, 100, 105]
+
+epoch_counter: !new:speechbrain.utils.epoch_loop.EpochCounter
+    limit: !ref <number_of_epochs>
+
+train_logger: !new:speechbrain.utils.train_logger.FileTrainLogger
+    save_file: !ref <train_log>
+
+ctc_stats: !name:speechbrain.utils.metric_stats.MetricStats
+    metric: !name:speechbrain.nnet.losses.ctc_loss
+        blank_index: !ref <blank_index>
+        reduction: batch
+
+per_stats: !name:speechbrain.utils.metric_stats.ErrorRateStats
+
+
+# # TIMIT
+# timit_local_data_folder: "/common/db/TIMIT"  # Path to TIMIT datase
+
+seed: 3047
+__set_seed: !apply:torch.manual_seed [!ref <seed>]
+
+# training parameters
+number_of_epochs: 100
+batch_size: 16
+lr: 0.0003
+sorting: ascending
+sample_rate: 16000
+gradient_accumulation: 2
+lr_pretrained: 0.00001
+
+# Mix-Precision Training
+auto_mix_prec: true
+# or
+precision: fp16         # 支持 "fp32"、"fp16" 或 "bf16"
+eval_precision: fp32    # 推理同样切换到 FP16
+
+# Dataloader options
+train_dataloader_opts:
+    batch_size: !ref <batch_size>
+    
+
+valid_dataloader_opts:
+    batch_size: !ref <batch_size>
+    
+
+test_dataloader_opts:
+    batch_size: 1
+
+pretrainer: !new:speechbrain.utils.parameter_transfer.Pretrainer
+    collect_in: !ref <save_folder>/
+    loadables:
+        perceived_ssl: !ref <perceived_ssl>
+        model: !ref <model>
+        tokenizer: !ref <tokenizer>
+
+encoder: !new:speechbrain.nnet.containers.LengthsCapableSequential
+    perceived_ssl: !ref <perceived_ssl>
+    enc: !ref <enc>
+    ctc_lin: !ref <ctc_lin>
+    log_softmax: !ref <log_softmax>
+
+decoding_function: !name:speechbrain.decoders.ctc_greedy_decode
+    blank_id: !ref <blank_index>
+
+tokenizer: !new:speechbrain.dataio.encoder.CTCTextEncoder
+
+modules:
+  encoder: !ref <encoder>